Add rs_sample.py

rs_sample.py

@@ -0,0 +1,108 @@
+"""RS sampling from C18-2 model (48.5% GSM8K)"""
+import json, re, asyncio, random
+from openai import AsyncOpenAI
+
+SP = "주어진 수학 문제를 단계별로 풀고 답변을 작성하세요.\n반드시 최종 답변을 \\boxed{정수} 형식으로 마지막 줄에 출력하세요.\n예시: \\boxed{42}"
+
+def extract_boxed(text):
+    m = re.findall(r'\\boxed\{([^}]+)\}', text)
+    return m[-1].strip() if m else None
+
+def normalize(a):
+    if a is None: return None
+    s = str(a).replace(",","").replace(" ","").strip()
+    try:
+        n = float(s)
+        return str(int(n)) if n == int(n) else str(n)
+    except: return s
+
+# Load GSM8K train questions (the ones we have gold answers for)
+with open("data/GSM8K_full_qwen3_30b.json") as f:
+    data = json.load(f)
+
+# Get unique questions with their gold answers
+q_to_gold = {}
+for d in data:
+    q = d["question"]
+    if q not in q_to_gold:
+        # Extract gold from the existing correct solutions
+        gold = extract_boxed(d["answer"])
+        if gold:
+            q_to_gold[q] = normalize(gold)
+
+questions = list(q_to_gold.keys())
+random.seed(42)
+random.shuffle(questions)
+# Sample a subset for RS (use all unique questions)
+print(f"Total unique questions: {len(questions)}")
+
+client = AsyncOpenAI(base_url="http://localhost:8000/v1", api_key="token-abc123")
+
+async def sample_one(question, n=16):
+    messages = [{"role": "user", "content": SP + "\n\n" + question}]
+    try:
+        resp = await client.chat.completions.create(
+            model="outputs/models/c18-2-combined-rs",
+            messages=messages, temperature=0.8, max_tokens=2048, n=n
+        )
+        return [c.message.content for c in resp.choices]
+    except Exception as e:
+        print(f"  Error: {e}")
+        return []
+
+async def main():
+    sem = asyncio.Semaphore(100)
+    sft_data = []
+    dpo_data = []
+    batch_size = 200
+
+    for batch_start in range(0, len(questions), batch_size):
+        batch = questions[batch_start:batch_start+batch_size]
+
+        async def process(q):
+            async with sem:
+                return q, await sample_one(q)
+
+        results = await asyncio.gather(*[process(q) for q in batch])
+
+        for q, answers in results:
+            gold = q_to_gold[q]
+            correct = []
+            incorrect = []
+            for a in answers:
+                pred = normalize(extract_boxed(a))
+                if pred and pred == gold:
+                    correct.append(a)
+                else:
+                    incorrect.append(a)
+
+            if correct:
+                best = min(correct, key=len)  # shortest correct
+                sft_data.append({
+                    "question": q, "answer": best,
+                    "n_correct": len(correct), "n_total": len(answers)
+                })
+
+            if correct and incorrect:
+                dpo_data.append({
+                    "question": q,
+                    "answer": min(correct, key=len),
+                    "bad_answer": max(incorrect, key=len)
+                })
+
+        print(f"  Batch {batch_start//batch_size + 1}: {len(sft_data)} sft, {len(dpo_data)} dpo")
+
+    import os
+    os.makedirs("outputs/c18_rs", exist_ok=True)
+    with open("outputs/c18_rs/sft_dataset.json", "w") as f:
+        json.dump(sft_data, f, ensure_ascii=False, indent=2)
+    with open("outputs/c18_rs/dpo_dataset.json", "w") as f:
+        json.dump(dpo_data, f, ensure_ascii=False, indent=2)
+
+    n4 = sum(1 for d in sft_data if d["n_correct"] >= 4)
+    print(f"\nRS Summary:")
+    print(f"  SFT: {len(sft_data)} (4+/16 filter: {n4})")
+    print(f"  DPO: {len(dpo_data)} pairs")
+    print(f"  Avg correct: {sum(d['n_correct'] for d in sft_data)/len(sft_data):.1f}/16")
+
+asyncio.run(main())