255 lines
10 KiB
Python
255 lines
10 KiB
Python
# /// script
|
||
# dependencies = [
|
||
# "openai>=1.0",
|
||
# "azure-identity",
|
||
# ]
|
||
# ///
|
||
"""
|
||
generate_distillation_data.py — Generate training data from a teacher model for distillation.
|
||
|
||
Creates a synthetic SFT dataset by:
|
||
1. Generating diverse prompts from combinatorial axes (topics × formats × contexts)
|
||
2. Having the teacher model produce responses
|
||
3. Quality-grading each response with an LLM judge
|
||
4. Filtering low-quality examples
|
||
5. Splitting into train/val/test JSONL files
|
||
|
||
Usage:
|
||
python generate_distillation_data.py \
|
||
--teacher gpt-4.1-mini \
|
||
--system-prompt "You are a formal business writer." \
|
||
--topics "earnings,risk,compliance" \
|
||
--num-prompts 300 \
|
||
--min-score 7.0 \
|
||
--output-dir ./my_dataset
|
||
|
||
# Or with a prompts file (one prompt per line):
|
||
python generate_distillation_data.py \
|
||
--teacher gpt-4.1-mini \
|
||
--prompts-file my_prompts.txt \
|
||
--output-dir ./my_dataset
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import random
|
||
import re
|
||
import sys
|
||
|
||
try:
|
||
sys.stdout.reconfigure(encoding="utf-8")
|
||
sys.stderr.reconfigure(encoding="utf-8")
|
||
except (AttributeError, OSError):
|
||
pass # Stream not reconfigurable (older Python or non-tty); default encoding is fine
|
||
import time
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
from common import HelpOnErrorParser, get_clients, _clamp_score
|
||
|
||
import openai
|
||
|
||
|
||
def verify_deployment(client, model):
|
||
"""Verify a model deployment exists by sending a trivial request."""
|
||
try:
|
||
client.chat.completions.create(
|
||
model=model,
|
||
messages=[{"role": "user", "content": "Hi"}],
|
||
max_completion_tokens=1,
|
||
)
|
||
return True
|
||
except openai.NotFoundError:
|
||
return False
|
||
except Exception:
|
||
return True # other errors (rate limit, etc.) mean the deployment exists
|
||
|
||
|
||
def generate_combinatorial_prompts(topics, formats, contexts, n):
|
||
"""Generate diverse prompts from combinatorial axes."""
|
||
prompts = []
|
||
for _ in range(n):
|
||
t = random.choice(topics)
|
||
f = random.choice(formats)
|
||
c = random.choice(contexts)
|
||
prompts.append(f"Context: {c}\n\nWrite {f} about: {t}.")
|
||
return prompts
|
||
|
||
|
||
def teacher_generate(client, model, system_prompt, prompt, retries=3):
|
||
"""Generate a single response from the teacher."""
|
||
for attempt in range(retries):
|
||
try:
|
||
resp = client.chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
temperature=0.7,
|
||
max_completion_tokens=1024,
|
||
)
|
||
return resp.choices[0].message.content
|
||
except Exception as e:
|
||
if attempt >= retries - 1:
|
||
print(f" Failed after {retries} attempts: {e}")
|
||
return None
|
||
time.sleep(2 * (attempt + 1))
|
||
return None
|
||
|
||
|
||
QUALITY_PROMPT = """Rate this AI-generated text on quality dimensions (1-10 each).
|
||
|
||
## Text to evaluate
|
||
{output}
|
||
|
||
## Dimensions
|
||
**Accuracy** (1-10): Is the content factually sound and coherent?
|
||
**Quality** (1-10): Is it well-written, clear, and professional?
|
||
**Task-fit** (1-10): Does it match the requested format and purpose?
|
||
|
||
Return ONLY JSON: {{"accuracy": <int>, "quality": <int>, "task_fit": <int>}}"""
|
||
|
||
|
||
def grade_output(client, judge_model, output, retries=3):
|
||
for attempt in range(retries):
|
||
try:
|
||
resp = client.chat.completions.create(
|
||
model=judge_model,
|
||
messages=[{"role": "user", "content": QUALITY_PROMPT.format(output=output)}],
|
||
temperature=0.0,
|
||
max_completion_tokens=100,
|
||
)
|
||
text = (resp.choices[0].message.content or "").strip()
|
||
match = re.search(r'\{[^}]+\}', text)
|
||
if match:
|
||
scores = json.loads(match.group())
|
||
return {k: _clamp_score(v) for k, v in scores.items()}
|
||
except Exception:
|
||
if attempt < retries - 1:
|
||
time.sleep(2)
|
||
return None
|
||
|
||
|
||
def main():
|
||
parser = HelpOnErrorParser(description="Generate distillation training data from a teacher model")
|
||
parser.add_argument("--base-url", default=os.environ.get("OPENAI_BASE_URL"),
|
||
help="Project /v1/ URL (preferred)")
|
||
parser.add_argument("--endpoint", default=os.environ.get("AZURE_OPENAI_ENDPOINT"),
|
||
help="Azure OpenAI endpoint (fallback)")
|
||
parser.add_argument("--project-endpoint", default=os.environ.get("AZURE_AI_PROJECT_ENDPOINT"),
|
||
help="Azure AI project endpoint (Foundry SDK)")
|
||
parser.add_argument("--api-key", default=os.environ.get("AZURE_OPENAI_API_KEY"))
|
||
parser.add_argument("--teacher", required=True, help="Teacher model deployment name")
|
||
parser.add_argument("--judge", default=None, help="Judge model (default: same as teacher)")
|
||
parser.add_argument("--system-prompt", default="You are a helpful assistant.", help="System prompt for teacher")
|
||
|
||
# Prompt generation (either combinatorial or from file)
|
||
parser.add_argument("--prompts-file", help="File with one prompt per line (skips combinatorial generation)")
|
||
parser.add_argument("--topics", help="Comma-separated topics for combinatorial prompts")
|
||
parser.add_argument("--formats", default="a concise response,a brief summary,a detailed explanation",
|
||
help="Comma-separated output formats")
|
||
parser.add_argument("--contexts", default="", help="Comma-separated context sentences")
|
||
parser.add_argument("--num-prompts", type=int, default=300, help="Number of prompts to generate")
|
||
|
||
# Quality
|
||
parser.add_argument("--min-score", type=float, default=7.0, help="Minimum average quality score to keep")
|
||
parser.add_argument("--skip-grading", action="store_true", help="Skip quality grading (keep all)")
|
||
|
||
# Output
|
||
parser.add_argument("--output-dir", default="./distillation_data", help="Output directory")
|
||
parser.add_argument("--train-split", type=float, default=0.8)
|
||
parser.add_argument("--val-split", type=float, default=0.1)
|
||
|
||
args = parser.parse_args()
|
||
|
||
client, method = get_clients(
|
||
base_url=args.base_url, azure_endpoint=args.endpoint,
|
||
project_endpoint=args.project_endpoint, api_key=args.api_key
|
||
)
|
||
judge = args.judge or args.teacher
|
||
|
||
# Step 0: Verify deployments exist
|
||
print(f"Verifying deployment '{args.teacher}'...")
|
||
if not verify_deployment(client, args.teacher):
|
||
print(f" ERROR: Deployment '{args.teacher}' not found. Available deployments can be listed in Azure Portal.")
|
||
sys.exit(1)
|
||
print(f" ✅ Teacher deployment verified.")
|
||
|
||
if judge != args.teacher:
|
||
print(f"Verifying judge deployment '{judge}'...")
|
||
if not verify_deployment(client, judge):
|
||
print(f" ERROR: Judge deployment '{judge}' not found.")
|
||
sys.exit(1)
|
||
print(f" ✅ Judge deployment verified.")
|
||
|
||
# Step 1: Generate or load prompts
|
||
if args.prompts_file:
|
||
with open(args.prompts_file, encoding="utf-8") as pf:
|
||
prompts = [line.strip() for line in pf if line.strip()]
|
||
print(f"Loaded {len(prompts)} prompts from {args.prompts_file}")
|
||
else:
|
||
topics = [t.strip() for t in (args.topics or "general knowledge").split(",")]
|
||
formats = [f.strip() for f in args.formats.split(",")]
|
||
contexts = [c.strip() for c in args.contexts.split(",") if c.strip()] or [""]
|
||
prompts = generate_combinatorial_prompts(topics, formats, contexts, args.num_prompts)
|
||
print(f"Generated {len(prompts)} prompts ({len(topics)} topics × {len(formats)} formats × {len(contexts)} contexts)")
|
||
|
||
# Step 2: Teacher generates responses
|
||
print(f"\nTeacher ({args.teacher}) generating responses...")
|
||
examples = []
|
||
for i, prompt in enumerate(prompts):
|
||
response = teacher_generate(client, args.teacher, args.system_prompt, prompt)
|
||
if response:
|
||
examples.append({"prompt": prompt, "response": response})
|
||
if (i + 1) % 25 == 0:
|
||
print(f" {i+1}/{len(prompts)} ({len(examples)} successful)")
|
||
print(f" Teacher produced {len(examples)}/{len(prompts)} responses")
|
||
|
||
# Step 3: Quality grade and filter
|
||
if not args.skip_grading:
|
||
print(f"\nGrading with {judge}...")
|
||
for i, ex in enumerate(examples):
|
||
scores = grade_output(client, judge, ex["response"])
|
||
if scores:
|
||
ex["scores"] = scores
|
||
ex["avg_score"] = sum(scores.values()) / len(scores)
|
||
else:
|
||
ex["avg_score"] = 0
|
||
if (i + 1) % 25 == 0:
|
||
print(f" Graded {i+1}/{len(examples)}")
|
||
|
||
filtered = [ex for ex in examples if ex["avg_score"] >= args.min_score]
|
||
avgs = [ex["avg_score"] for ex in examples if ex["avg_score"] > 0]
|
||
print(f" Passed filter (>= {args.min_score}): {len(filtered)}/{len(examples)}")
|
||
if avgs:
|
||
print(f" Scores: min={min(avgs):.1f}, max={max(avgs):.1f}, mean={sum(avgs)/len(avgs):.1f}")
|
||
else:
|
||
filtered = examples
|
||
print(f"Skipping grading — keeping all {len(filtered)} examples")
|
||
|
||
# Step 4: Convert to SFT format and split
|
||
sft_data = [{"messages": [
|
||
{"role": "system", "content": args.system_prompt},
|
||
{"role": "user", "content": ex["prompt"]},
|
||
{"role": "assistant", "content": ex["response"]},
|
||
]} for ex in filtered]
|
||
|
||
random.shuffle(sft_data)
|
||
n = len(sft_data)
|
||
t_end = int(n * args.train_split)
|
||
v_end = int(n * (args.train_split + args.val_split))
|
||
splits = {"train": sft_data[:t_end], "validation": sft_data[t_end:v_end], "test": sft_data[v_end:]}
|
||
|
||
os.makedirs(args.output_dir, exist_ok=True)
|
||
for name, data in splits.items():
|
||
path = os.path.join(args.output_dir, f"{name}.jsonl")
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
for ex in data:
|
||
f.write(json.dumps(ex, ensure_ascii=False) + "\n")
|
||
print(f" {name}: {len(data)} examples → {path}")
|
||
|
||
print(f"\n✅ Done! Dataset ready in {args.output_dir}/")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|