File size: 3,371 Bytes
d7cc8ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | # /// script
# dependencies = ["datasets", "huggingface_hub>=0.34"]
# ///
"""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()
|