tinymind-native-8b-remote-handoff/bundle/data/distill_claude.py

"""
Claude API Distillation — สร้าง high-quality reasoning traces

ใช้ Claude เป็น teacher model สร้าง:
1. Chain-of-Thought QA pairs (มี <think> traces)
2. DPO preference pairs (chosen=มี reasoning, rejected=ตอบตรงๆ)
3. Math/logic reasoning problems พร้อม step-by-step solution
4. Diverse instruction-following examples

ต้องตั้งค่า ANTHROPIC_API_KEY ก่อน
"""

from __future__ import annotations

import json
import os
import re
import time
from pathlib import Path
import anthropic

# ─── Config ──────────────────────────────────────────────────────────────────

OUTPUT_DIR = Path(__file__).parent / "filtered"
OUTPUT_DIR.mkdir(exist_ok=True)

COT_OUTPUT       = OUTPUT_DIR / "cot_qa.jsonl"
DPO_OUTPUT       = OUTPUT_DIR / "dpo_pairs.jsonl"
REASONING_OUTPUT = OUTPUT_DIR / "reasoning_qa.jsonl"

MODEL = "claude-opus-4-7"   # ใช้โมเดลดีที่สุดเป็น teacher
MAX_TOKENS = 1024
RATE_LIMIT_DELAY = 0.5      # วินาทีระหว่าง requests

# ─── Topics ──────────────────────────────────────────────────────────────────

THAI_REASONING_TOPICS = [
    "คณิตศาสตร์ระดับมัธยม", "ตรรกศาสตร์และการอนุมาน",
    "ฟิสิกส์พื้นฐาน", "เคมีพื้นฐาน",
    "การแก้ปัญหาเชิงตรรกะ", "สถิติและความน่าจะเป็น",
    "การเขียนโปรแกรม Python พื้นฐาน", "ชีววิทยาและชีวเคมี",
]

EN_REASONING_TOPICS = [
    "algebra and calculus", "formal logic and proofs",
    "physics problems", "algorithm design",
    "probability and statistics", "number theory",
    "combinatorics", "geometry proofs",
]

THAI_KNOWLEDGE_TOPICS = [
    # วิชาการ
    "ประวัติศาสตร์ไทยและอาเซียน", "วิทยาศาสตร์และเทคโนโลยี",
    "เศรษฐกิจและการเงิน", "วัฒนธรรมและสังคม",
    "สุขภาพและการแพทย์", "กฎหมายและจริยธรรม",
    "สิ่งแวดล้อมและพลังงาน", "AI และวิทยาการคอมพิวเตอร์",
    # ONET/PAT ระดับ
    "คณิตศาสตร์ ม.ปลาย (ONET)", "ฟิสิกส์ระดับมหาวิทยาลัยปี 1",
    "เคมีระดับมหาวิทยาลัยปี 1", "ชีววิทยาระดับมหาวิทยาลัยปี 1",
    "ภาษาไทยระดับสูง: วรรณคดีและการวิเคราะห์",
    "สังคมศึกษา ม.ปลาย: การเมืองการปกครอง",
    # เฉพาะทาง
    "ปรัชญาพุทธศาสนาและจริยศาสตร์ไทย",
    "ภูมิปัญญาไทย นวัตกรรมและการประยุกต์ใช้",
]

EN_KNOWLEDGE_TOPICS = [
    # MMLU categories (all 57 mapped to groups)
    "abstract algebra and number theory",
    "anatomy and physiology",
    "astronomy and astrophysics",
    "business ethics and corporate governance",
    "clinical knowledge and medical diagnosis",
    "college chemistry and organic reactions",
    "college computer science and algorithms",
    "college mathematics and proof-writing",
    "college physics and quantum mechanics",
    "college biology and molecular biology",
    "conceptual physics and mechanics",
    "econometrics and statistical inference",
    "electrical engineering and circuits",
    "elementary mathematics and problem solving",
    "formal logic and propositional calculus",
    "global facts and world geography",
    "high school biology and genetics",
    "high school chemistry and stoichiometry",
    "high school computer science and data structures",
    "high school government and political science",
    "high school macroeconomics",
    "high school mathematics functions and graphs",
    "high school physics kinematics and dynamics",
    "high school psychology and cognitive science",
    "high school statistics and probability",
    "high school world history modern era",
    "human aging and gerontology",
    "international law and treaties",
    "jurisprudence and legal philosophy",
    "machine learning and deep learning concepts",
    "management and organizational behavior",
    "marketing and consumer behavior",
    "medical genetics and hereditary disease",
    "miscellaneous trivia and general knowledge",
    "moral philosophy and applied ethics",
    "nutrition and dietetics",
    "philosophy of mind and consciousness",
    "prehistory and ancient civilizations",
    "professional accounting and auditing",
    "professional law and bar exam",
    "professional medicine clinical reasoning",
    "professional psychology DSM diagnosis",
    "public relations and crisis communication",
    "security studies and geopolitics",
    "sociology and social stratification",
    "us foreign policy and diplomacy",
    "virology and infectious disease",
    "world religions and comparative theology",
]


# ─── Prompt Templates ────────────────────────────────────────────────────────

def cot_prompt_thai(topic: str, n: int = 5) -> str:
    return f"""สร้างชุดคำถาม-คำตอบ {n} ข้อเกี่ยวกับ "{topic}"

สำหรับแต่ละข้อ ให้:
1. ตั้งคำถามที่ต้องใช้การคิดวิเคราะห์
2. แสดงกระบวนการคิดอย่างละเอียดใน <think>...</think>
3. ให้คำตอบสุดท้ายที่ชัดเจนใน <answer>...</answer>

ตอบเป็น JSON array:
[
  {{
    "question": "คำถาม",
    "thinking": "กระบวนการคิดทีละขั้น...",
    "answer": "คำตอบสุดท้าย",
    "lang": "th",
    "topic": "{topic}"
  }}
]

กฎ: ถูกต้อง 100% ไม่แต่งเรื่อง ตอบ JSON เท่านั้น"""


def cot_prompt_en(topic: str, n: int = 5) -> str:
    return f"""Generate {n} Q&A pairs about "{topic}" that require analytical thinking.

For each pair:
1. Ask a question requiring multi-step reasoning
2. Show detailed reasoning inside <think>...</think>
3. Give a clear final answer inside <answer>...</answer>

Return JSON array:
[
  {{
    "question": "...",
    "thinking": "step by step reasoning...",
    "answer": "final answer",
    "lang": "en",
    "topic": "{topic}"
  }}
]

Rules: 100% factually correct, no hallucination, JSON only."""


def dpo_prompt_thai(topic: str, n: int = 3) -> str:
    return f"""สร้างชุด preference pairs {n} ชุดเกี่ยวกับ "{topic}"

แต่ละชุดมี:
- question: คำถาม
- chosen: คำตอบที่ดี (มีการคิดวิเคราะห์ใน <think>...</think> และคำตอบใน <answer>...</answer>)
- rejected: คำตอบที่แย่กว่า (ตอบตรงๆ สั้นเกิน หรือไม่ครบถ้วน)

ตอบเป็น JSON array:
[
  {{
    "question": "คำถาม",
    "chosen": "<think>\\nคิดวิเคราะห์...\\n</think>\\n<answer>คำตอบดี</answer>",
    "rejected": "คำตอบสั้นๆ ไม่ครบถ้วน",
    "lang": "th"
  }}
]"""


def dpo_prompt_en(topic: str, n: int = 3) -> str:
    return f"""Generate {n} preference pairs about "{topic}".

Each pair:
- question: the question
- chosen: good answer (with <think>...</think> reasoning + <answer>...</answer>)
- rejected: worse answer (too brief, no reasoning, or incomplete)

Return JSON array:
[
  {{
    "question": "...",
    "chosen": "<think>\\nDetailed reasoning...\\n</think>\\n<answer>Good answer</answer>",
    "rejected": "Short incomplete answer",
    "lang": "en"
  }}
]"""


def math_reasoning_prompt(lang: str, n: int = 5) -> str:
    if lang == "th":
        return f"""สร้างโจทย์คณิตศาสตร์/ตรรกะ {n} ข้อพร้อมคำเฉลย

แต่ละข้อต้องมี:
- question: โจทย์ที่ชัดเจน
- thinking: วิธีคิดทีละขั้นตอนอย่างละเอียด
- answer: คำตอบสุดท้าย (ตัวเลขหรือข้อความที่ชัดเจน)
- ground_truth: คำตอบที่ verify ได้ (ตัวเลขหรือ true/false)

ตอบ JSON array เท่านั้น ตัวอย่าง:
[{{"question": "ถ้า 3x + 7 = 22 แล้ว x = ?",
  "thinking": "นำ 7 ออก: 3x = 15\\nหาร 3: x = 5",
  "answer": "x = 5",
  "ground_truth": "5",
  "lang": "th", "topic": "algebra"}}]"""
    else:
        return f"""Generate {n} math/logic problems with full solutions.

Each must have:
- question: clear problem statement
- thinking: detailed step-by-step solution
- answer: clear final answer
- ground_truth: verifiable answer (number or true/false)

Return JSON array only. Example:
[{{"question": "If 3x + 7 = 22, what is x?",
  "thinking": "Subtract 7: 3x = 15\\nDivide by 3: x = 5",
  "answer": "x = 5",
  "ground_truth": "5",
  "lang": "en", "topic": "algebra"}}]"""


# ─── Claude API ───────────────────────────────────────────────────────────────

def call_claude(client: anthropic.Anthropic, prompt: str) -> str:
    try:
        msg = client.messages.create(
            model=MODEL,
            max_tokens=MAX_TOKENS,
            messages=[{"role": "user", "content": prompt}],
            system=(
                "You are a high-quality training data generator. "
                "Always output valid JSON only. No markdown, no explanation."
            ),
        )
        block = msg.content[0] if msg.content else None
        return getattr(block, "text", "") or ""
    except anthropic.RateLimitError:
        print("  Rate limit — waiting 60s ...")
        time.sleep(60)
        return call_claude(client, prompt)
    except Exception as e:
        print(f"  API error: {e}")
        return ""


def extract_json_list(text: str) -> list[dict]:
    text = text.strip()
    match = re.search(r"\[[\s\S]*\]", text)
    if not match:
        return []
    try:
        data = json.loads(match.group())
        return data if isinstance(data, list) else []
    except Exception:
        return []


# ─── Generation Pipelines ─────────────────────────────────────────────────────

def generate_cot_data(
    client: anthropic.Anthropic,
    output_path: Path,
    rounds_per_topic: int = 2,
) -> int:
    total = 0
    with open(output_path, "a", encoding="utf-8") as f:
        for topic in THAI_KNOWLEDGE_TOPICS:
            for _ in range(rounds_per_topic):
                raw = call_claude(client, cot_prompt_thai(topic))
                pairs = extract_json_list(raw)
                for p in pairs:
                    if p.get("question") and p.get("thinking") and p.get("answer"):
                        p["source"] = "claude_cot"
                        p["context"] = ""
                        f.write(json.dumps(p, ensure_ascii=False) + "\n")
                        total += 1
                time.sleep(RATE_LIMIT_DELAY)
            print(f"  Thai CoT [{topic}]: {total} total")

        for topic in EN_KNOWLEDGE_TOPICS:
            for _ in range(rounds_per_topic):
                raw = call_claude(client, cot_prompt_en(topic))
                pairs = extract_json_list(raw)
                for p in pairs:
                    if p.get("question") and p.get("thinking") and p.get("answer"):
                        p["source"] = "claude_cot"
                        p["context"] = ""
                        f.write(json.dumps(p, ensure_ascii=False) + "\n")
                        total += 1
                time.sleep(RATE_LIMIT_DELAY)
            print(f"  EN CoT [{topic}]: {total} total")

    return total


def generate_dpo_data(
    client: anthropic.Anthropic,
    output_path: Path,
    rounds_per_topic: int = 2,
) -> int:
    total = 0
    all_topics = [
        (t, "th") for t in THAI_KNOWLEDGE_TOPICS
    ] + [
        (t, "en") for t in EN_KNOWLEDGE_TOPICS
    ]

    with open(output_path, "a", encoding="utf-8") as f:
        for topic, lang in all_topics:
            for _ in range(rounds_per_topic):
                if lang == "th":
                    prompt = dpo_prompt_thai(topic)
                else:
                    prompt = dpo_prompt_en(topic)
                raw = call_claude(client, prompt)
                pairs = extract_json_list(raw)
                for p in pairs:
                    if p.get("question") and p.get("chosen") and p.get("rejected"):
                        p["source"] = "claude_dpo"
                        f.write(json.dumps(p, ensure_ascii=False) + "\n")
                        total += 1
                time.sleep(RATE_LIMIT_DELAY)
        print(f"  DPO pairs: {total} total")

    return total


def generate_reasoning_data(
    client: anthropic.Anthropic,
    output_path: Path,
    rounds: int = 5,
) -> int:
    total = 0
    with open(output_path, "a", encoding="utf-8") as f:
        for lang in ["th", "en"]:
            topics = THAI_REASONING_TOPICS if lang == "th" else EN_REASONING_TOPICS
            for topic in topics:
                for _ in range(rounds):
                    prompt = math_reasoning_prompt(lang)
                    raw = call_claude(client, prompt)
                    pairs = extract_json_list(raw)
                    for p in pairs:
                        if p.get("question") and p.get("answer"):
                            p["source"] = "claude_reasoning"
                            p["topic"] = topic
                            f.write(json.dumps(p, ensure_ascii=False) + "\n")
                            total += 1
                    time.sleep(RATE_LIMIT_DELAY)
        print(f"  Reasoning: {total} total")
    return total


# ─── MMLU Deep Distillation ───────────────────────────────────────────────────

MMLU_OUTPUT = OUTPUT_DIR / "mmlu_cot.jsonl"

def mmlu_prompt(category: str, n: int = 8) -> str:
    return f"""Generate {n} MMLU-style multiple-choice questions about "{category}".

Each question must:
1. Have exactly 4 choices labeled A, B, C, D
2. Have one unambiguously correct answer
3. Include step-by-step reasoning explaining WHY the answer is correct
4. Be at university/professional level difficulty

Return JSON array only:
[
  {{
    "question": "...",
    "choices": {{"A": "...", "B": "...", "C": "...", "D": "..."}},
    "correct": "B",
    "thinking": "Step-by-step: First consider... The answer is B because...",
    "answer": "<think>\\n...reasoning...\\n</think>\\n<answer>B</answer>",
    "category": "{category}",
    "lang": "en"
  }}
]"""


def generate_mmlu_data(
    client: anthropic.Anthropic,
    output_path: Path = MMLU_OUTPUT,
    rounds_per_category: int = 2,
) -> int:
    total = 0
    with open(output_path, "a", encoding="utf-8") as f:
        for category in EN_KNOWLEDGE_TOPICS:
            for _ in range(rounds_per_category):
                raw = call_claude(client, mmlu_prompt(category))
                items = extract_json_list(raw)
                for item in items:
                    if item.get("question") and item.get("correct") and item.get("thinking"):
                        item["source"] = "claude_mmlu"
                        item["topic"] = "mmlu"
                        item["context"] = ""
                        f.write(json.dumps(item, ensure_ascii=False) + "\n")
                        total += 1
                time.sleep(RATE_LIMIT_DELAY)
        print(f"  MMLU: {total} questions")
    return total


# ─── Code Distillation ────────────────────────────────────────────────────────

CODE_COT_OUTPUT = OUTPUT_DIR / "code_cot.jsonl"

CODE_CATEGORIES = [
    "Python data manipulation with pandas and numpy",
    "object-oriented programming design patterns",
    "recursive algorithms and divide-and-conquer",
    "dynamic programming optimization problems",
    "graph algorithms: Dijkstra, A*, Floyd-Warshall",
    "tree traversal and binary search trees",
    "system design: REST API, database schema",
    "concurrency and async programming in Python",
    "regular expressions and text parsing",
    "competitive programming: prefix sums, segment trees",
]

def code_cot_prompt(category: str, n: int = 4) -> str:
    return f"""Generate {n} coding problems about "{category}" with complete solutions.

Each problem must have:
- question: clear problem statement with function signature
- thinking: step-by-step algorithm design with complexity analysis
- answer: complete, correct, well-commented Python code
- test_cases: 3+ pytest-style assertions
- time_complexity: Big-O notation
- ground_truth: expected output for first test case

Return JSON array:
[
  {{
    "question": "Write def func(args) -> type that ...",
    "thinking": "Algorithm: 1. ... 2. ... Time: O(...) Space: O(...)",
    "answer": "def func(args):\\n    ...complete implementation...",
    "test_cases": "assert func(x) == y\\nassert func(z) == w",
    "time_complexity": "O(n log n)",
    "ground_truth": "expected_output",
    "category": "{category}",
    "lang": "en",
    "source": "claude_code",
    "topic": "programming"
  }}
]"""


def generate_code_cot_data(
    client: anthropic.Anthropic,
    output_path: Path = CODE_COT_OUTPUT,
    rounds_per_category: int = 3,
) -> int:
    total = 0
    with open(output_path, "a", encoding="utf-8") as f:
        for category in CODE_CATEGORIES:
            for _ in range(rounds_per_category):
                raw = call_claude(client, code_cot_prompt(category))
                items = extract_json_list(raw)
                for item in items:
                    if item.get("question") and item.get("answer") and item.get("test_cases"):
                        item["context"] = ""
                        f.write(json.dumps(item, ensure_ascii=False) + "\n")
                        total += 1
                time.sleep(RATE_LIMIT_DELAY)
        print(f"  Code CoT: {total} problems")
    return total


# ─── Thai Deep Distillation ───────────────────────────────────────────────────

THAI_DEEP_OUTPUT = OUTPUT_DIR / "thai_deep_cot.jsonl"

THAI_DEEP_TOPICS = [
    ("ONET คณิตศาสตร์ ม.6: ฟังก์ชันและกราฟ", "th"),
    ("PAT1 คณิตศาสตร์: แคลคูลัสเบื้องต้น", "th"),
    ("ONET ฟิสิกส์: กลศาสตร์และอุณหพลศาสตร์", "th"),
    ("ONET เคมี: ปฏิกิริยาเคมีและสมดุล", "th"),
    ("ONET ชีววิทยา: พันธุศาสตร์และวิวัฒนาการ", "th"),
    ("ภาษาไทยชั้นสูง: การวิเคราะห์วรรณกรรมไทย", "th"),
    ("ประวัติศาสตร์ไทย: สมัยรัตนโกสินทร์", "th"),
    ("เศรษฐศาสตร์ไทย: นโยบายการคลังและการเงิน", "th"),
    ("กฎหมายแพ่งและพาณิชย์ไทยเบื้องต้น", "th"),
    ("วัฒนธรรมและประเพณีไทยภูมิภาค", "th"),
    ("ภาษาไทย: ไวยากรณ์และการใช้ภาษา", "th"),
    ("พระพุทธศาสนา: หลักธรรมและการประยุกต์ใช้", "th"),
]

def thai_deep_prompt(topic: str, n: int = 6) -> str:
    return f"""สร้างชุดคำถาม-คำตอบระดับสูง {n} ข้อเกี่ยวกับ "{topic}"

แต่ละข้อต้องมีระดับความยากเทียบกับ ONET/GATPAT หรือระดับมหาวิทยาลัยปี 1

รูปแบบ JSON:
[
  {{
    "question": "คำถามที่ต้องคิดวิเคราะห์เชิงลึก",
    "thinking": "การวิเคราะห์ทีละขั้น ครบถ้วน และถูกต้อง",
    "answer": "<think>\\n...การวิเคราะห์...\\n</think>\\n<answer>คำตอบที่สมบูรณ์</answer>",
    "difficulty": "medium|hard",
    "lang": "th",
    "topic": "{topic}",
    "source": "claude_thai_deep"
  }}
]

ข้อกำหนด: ถูกต้อง 100% ระดับมาตรฐานสอบ ตอบ JSON เท่านั้น"""


def generate_thai_deep_data(
    client: anthropic.Anthropic,
    output_path: Path = THAI_DEEP_OUTPUT,
    rounds_per_topic: int = 3,
) -> int:
    total = 0
    with open(output_path, "a", encoding="utf-8") as f:
        for topic, lang in THAI_DEEP_TOPICS:
            for _ in range(rounds_per_topic):
                raw = call_claude(client, thai_deep_prompt(topic))
                items = extract_json_list(raw)
                for item in items:
                    if item.get("question") and item.get("thinking") and item.get("answer"):
                        item["context"] = ""
                        f.write(json.dumps(item, ensure_ascii=False) + "\n")
                        total += 1
                time.sleep(RATE_LIMIT_DELAY)
        print(f"  Thai deep: {total} questions")
    return total


# ─── Main ─────────────────────────────────────────────────────────────────────

def distill_with_claude(
    cot: bool = True,
    dpo: bool = True,
    reasoning: bool = True,
    mmlu: bool = True,
    code: bool = True,
    thai_deep: bool = True,
    rounds_per_topic: int = 2,
):
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        raise EnvironmentError(
            "ตั้งค่า ANTHROPIC_API_KEY ก่อน:\n"
            "  $env:ANTHROPIC_API_KEY = 'sk-ant-...'"
        )

    client = anthropic.Anthropic(api_key=api_key)
    print(f"Claude distillation | model={MODEL}\n")

    total = 0
    if cot:
        print("=== CoT QA pairs ===")
        n = generate_cot_data(client, COT_OUTPUT, rounds_per_topic)
        print(f"CoT: {n:,} pairs → {COT_OUTPUT}\n")
        total += n

    if dpo:
        print("=== DPO preference pairs ===")
        n = generate_dpo_data(client, DPO_OUTPUT, rounds_per_topic)
        print(f"DPO: {n:,} pairs → {DPO_OUTPUT}\n")
        total += n

    if reasoning:
        print("=== Math/Logic reasoning ===")
        n = generate_reasoning_data(client, REASONING_OUTPUT, rounds=rounds_per_topic)
        print(f"Reasoning: {n:,} pairs → {REASONING_OUTPUT}\n")
        total += n

    if mmlu:
        print("=== MMLU deep knowledge (all 48 categories) ===")
        n = generate_mmlu_data(client, MMLU_OUTPUT, rounds_per_category=rounds_per_topic)
        print(f"MMLU: {n:,} questions → {MMLU_OUTPUT}\n")
        total += n

    if code:
        print("=== Code CoT (10 categories) ===")
        n = generate_code_cot_data(client, CODE_COT_OUTPUT, rounds_per_category=rounds_per_topic)
        print(f"Code: {n:,} problems → {CODE_COT_OUTPUT}\n")
        total += n

    if thai_deep:
        print("=== Thai Deep (ONET/PAT level) ===")
        n = generate_thai_deep_data(client, THAI_DEEP_OUTPUT, rounds_per_topic=rounds_per_topic)
        print(f"Thai deep: {n:,} questions → {THAI_DEEP_OUTPUT}\n")
        total += n

    print(f"\n{'='*50}")
    print(f"Distillation complete: {total:,} total examples")
    print(f"Files in {OUTPUT_DIR}:")
    for f in sorted(OUTPUT_DIR.glob("*.jsonl")):
        lines = sum(1 for _ in open(f, encoding="utf-8"))
        print(f"  {f.name}: {lines:,} examples")


if __name__ == "__main__":
    distill_with_claude(rounds_per_topic=3)