1 files changed, 60 insertions, 0 deletions
diff --git a/data/templates.py b/data/templates.py
new file mode 100644
index 0000000..a4bc9a4
--- /dev/null
+++ b/data/templates.py
@@ -0,0 +1,60 @@
+"""Prompt templates for LongLaMP tasks."""
+
+
+SYSTEM_PROMPT = "You are a helpful writing assistant."
+
+
+def build_query_prompt(query_input: str, task: str) -> str:
+    """Build the inference prompt for a query (no personalization text)."""
+    if task == "review":
+        return (
+            f"{query_input}\n\n"
+            f"Write the review text:"
+        )
+    else:  # topic
+        return (
+            f"{query_input}\n\n"
+            f"Write the post content:"
+        )
+
+
+def build_support_prompt(support_input: str, task: str) -> str:
+    """Build the prompt for a support item (used in teacher forcing to cache hidden states)."""
+    if task == "review":
+        return (
+            f"{support_input}\n\n"
+            f"Write the review text:"
+        )
+    else:  # topic
+        return (
+            f"{support_input}\n\n"
+            f"Write the post content:"
+        )
+
+
+def build_prompt_with_examples(query_input: str, support_items: list, task: str) -> str:
+    """Build prompt with K support examples included as in-context demonstrations.
+
+    Used by Prompt-All-K and BM25-Top1 baselines.
+    """
+    parts = []
+    parts.append("Here are some examples of this user's writing style:\n")
+
+    for i, item in enumerate(support_items, 1):
+        parts.append(f"--- Example {i} ---")
+        parts.append(f"Prompt: {item['support_input']}")
+        parts.append(f"Response: {item['support_output']}")
+        parts.append("")
+
+    parts.append("Now, write in the same style as the examples above.\n")
+    parts.append(build_query_prompt(query_input, task))
+
+    return "\n".join(parts)
+
+
+def build_chat_messages(prompt: str) -> list:
+    """Wrap a prompt into chat message format for Qwen."""
+    return [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": prompt},
+    ]