Update README

README.md

@@ -85,6 +85,8 @@ GRPO is supervised with a sum of seven rewards, grouped into three families:
 
 ## Quickstart
 
+DecomposeRL expects a specific verification prompt around your `claim` + `evidence_doc`. The `build_prompt` helper below wraps them for you so you don't have to construct the full instruction block every time.
+
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -97,16 +99,8 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto",
 )
 
-evidence_doc = (
-    "The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, "
-    "France. It is named after the engineer Gustave Eiffel, whose company designed and "
-    "built the tower from 1887 to 1889. Locally nicknamed 'La dame de fer', it was "
-    "constructed as the centerpiece of the 1889 World's Fair. The tower is 330 metres "
-    "(1,083 ft) tall."
-)
-claim = "The Eiffel Tower was completed in 1887 and stands 330 metres tall."
 
-user_prompt = f"""You are tasked with systematically verifying the accuracy of a claim. You will be provided with a claim to verify and an evidence document to consult.
+PROMPT_TEMPLATE = """You are tasked with systematically verifying the accuracy of a claim. You will be provided with a claim to verify and an evidence document to consult.
 
 Here is the evidence document you should consult:
 
@@ -130,13 +124,37 @@ Stop immediately after the closing </verification> tag.
 
 Begin your verification process now."""
 
-messages = [{"role": "user", "content": user_prompt}]
-text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-# max_new_tokens matches training-time max_completion_length
-out = model.generate(**inputs, max_new_tokens=4500, temperature=0.7, do_sample=True)
-response = tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+def build_prompt(claim: str, evidence_doc: str) -> str:
+    """Wrap a claim + evidence document in the DecomposeRL verification prompt."""
+    return PROMPT_TEMPLATE.format(claim=claim, evidence_doc=evidence_doc)
+
+
+def verify(claim: str, evidence_doc: str, max_new_tokens: int = 4500, temperature: float = 0.7) -> str:
+    """Run the model end-to-end on a (claim, evidence_doc) pair and return the raw trace."""
+    messages = [{"role": "user", "content": build_prompt(claim, evidence_doc)}]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    out = model.generate(
+        **inputs,
+        max_new_tokens=max_new_tokens,  # matches training-time max_completion_length
+        temperature=temperature,
+        do_sample=True,
+    )
+    return tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+
+
+# Usage
+evidence_doc = (
+    "The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, "
+    "France. It is named after the engineer Gustave Eiffel, whose company designed and "
+    "built the tower from 1887 to 1889. Locally nicknamed 'La dame de fer', it was "
+    "constructed as the centerpiece of the 1889 World's Fair. The tower is 330 metres "
+    "(1,083 ft) tall."
+)
+claim = "The Eiffel Tower was completed in 1887 and stands 330 metres tall."
+
+response = verify(claim, evidence_doc)
 print(response)
 ```
 
@@ -154,9 +172,20 @@ def parse_trace(text: str):
     return [(tag, body.strip()) for tag, body in TAG_RE.findall(text)]
 
 def pretty_print(text: str) -> None:
+    parsed = parse_trace(text)
+    tags = {tag for tag, _ in parsed}
+    if not parsed or "verification" not in tags:
+        print("⚠️  Could not parse output into the expected "
+              "think/question/answer/verification structure.")
+        print("Raw output:")
+        print("─" * 78)
+        print(text)
+        print("─" * 78)
+        return
+
     cycle_idx = 0
     pending_q = None
-    for tag, body in parse_trace(text):
+    for tag, body in parsed:
         if tag == "think":
             print("─" * 78)
             print("🧠  THINK")