| |
| |
| |
| """Generate SDPO hints for the MATH dataset (hard levels) with a large model on an HF Inference Endpoint. |
| |
| Pushes two datasets built from the SAME problems: |
| - <output_repo> : problem, answer, solution, privileged_context (for A hints / B baseline) |
| - <output_repo>-plain : problem, answer, solution (for C gold_feedback) |
| |
| The large model only *generates* text (no logprobs), so any standard IE chat deployment works. |
| |
| Usage: |
| HF_TOKEN=hf_... python generate_hints_math.py \ |
| --endpoint_url https://xxxx.endpoints.huggingface.cloud \ |
| --output_repo sergiopaniego/math-sdpo-hints \ |
| --num_examples 600 --min_level 4 |
| """ |
| import argparse |
| import os |
| import time |
| from concurrent.futures import ThreadPoolExecutor |
|
|
| from datasets import load_dataset |
| from huggingface_hub import InferenceClient |
|
|
|
|
| HINT_SYSTEM = ( |
| "You are a math expert. Given a competition math problem, write a concise worked solution " |
| "(the key steps, not verbose) and end with the final answer as \\boxed{<answer>}." |
| ) |
|
|
|
|
| def make_hint(client, problem, model, max_retries=3): |
| for attempt in range(max_retries): |
| try: |
| resp = client.chat_completion( |
| messages=[ |
| {"role": "system", "content": HINT_SYSTEM}, |
| {"role": "user", "content": problem}, |
| ], |
| model=model, |
| max_tokens=512, |
| temperature=0.3, |
| ) |
| return resp.choices[0].message.content.strip() |
| except Exception: |
| if attempt == max_retries - 1: |
| return "" |
| time.sleep(2 * (attempt + 1)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--endpoint_url", required=True, help="IE endpoint base URL") |
| ap.add_argument("--model", default=None, help="Model id for chat_completion (optional on a dedicated endpoint)") |
| ap.add_argument("--num_examples", type=int, default=600) |
| ap.add_argument("--min_level", type=int, default=4, help="Keep only problems with level >= this (1-5).") |
| ap.add_argument("--output_repo", required=True, help="Hub dataset repo id (hints version).") |
| ap.add_argument("--concurrency", type=int, default=8) |
| args = ap.parse_args() |
|
|
| client = InferenceClient(base_url=args.endpoint_url, token=os.environ.get("HF_TOKEN")) |
|
|
| ds = load_dataset("nlile/hendrycks-MATH-benchmark", split="train") |
| ds = ds.filter(lambda ex: ex["level"] >= args.min_level) |
| ds = ds.shuffle(seed=42).select(range(min(args.num_examples, len(ds)))) |
| print(f"Selected {len(ds)} problems (level >= {args.min_level}).") |
|
|
| t0 = time.time() |
| with ThreadPoolExecutor(max_workers=args.concurrency) as pool: |
| hints = list(pool.map(lambda ex: make_hint(client, ex["problem"], args.model), ds)) |
|
|
| n_empty = sum(1 for h in hints if not h) |
| ds = ds.add_column("privileged_context", hints).filter(lambda ex: bool(ex["privileged_context"])) |
|
|
| ds.push_to_hub(args.output_repo) |
| ds.remove_columns("privileged_context").push_to_hub(f"{args.output_repo}-plain") |
| print(f"Pushed {len(ds)} examples to {args.output_repo} (+ -plain) ({n_empty} empty dropped) in {time.time() - t0:.0f}s") |
| print("Columns:", ds.column_names) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
