Add rs_sample.py

1dcce72 verified 8 months ago

3.72 kB

"""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())

	"""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())