"""A minimal Distilabel pipeline that generates and judges with your own local models. Purpose: the same generate-then-judge shape as generate-teacher-data.py plus judge.py, written as a Distilabel pipeline instead of as two scripts, so you can see what the framework buys you and what it costs. Two steps: a teacher answers every seed prompt, and a second model grades each answer against a rubric. Both point at an OpenAI-compatible endpoint on your own machine, so nothing leaves the house. Platform: all (pure Python over HTTP; the models may be served by any engine on any track, or by the Part 9 gateway under two aliases) Minimum memory: 12 GB on the machine serving the models; the pipeline process itself is small Assumes: Python 3.10 or newer and `pip install distilabel[openai]`. Distilabel is not one of the course's pinned tools, so check the version you installed against the documentation at https://distilabel.argilla.io/latest/ before relying on any behaviour here. A reachable OpenAI-compatible endpoint at --base-url. The seed file is JSON Lines with an "id" and a "prompt" on every line, as written by make-seed-prompts.py. Usage: python3 distilabel-pipeline.py --seeds seeds/prompts.jsonl \\ --base-url http://127.0.0.1:4000/v1 --teacher local/chat \\ --judge local/chat --out-dir distilabel-out --limit 40 python3 distilabel-pipeline.py --seeds seeds/prompts.jsonl --print-only --print-only writes nothing and starts no model; it prints the pipeline definition so the shape can be read on a machine with nothing installed. """ from __future__ import annotations import argparse import json import os from pathlib import Path # The judge prompt is a Jinja2 template, which is what TextGeneration's `template` # attribute takes. `columns` names the input columns the template may refer to. JUDGE_TEMPLATE = """You grade one answer against a rubric. Reply with JSON only, of the form {"score": <1-5>, "reason": ""}. Rubric: the answer must do what the instruction asked, in the shape it asked for, without adding facts the instruction did not supply. Instruction: {{ instruction }} Answer: {{ generation }} """ def build_pipeline(args: argparse.Namespace): """Assemble the pipeline. Imported lazily so --print-only needs nothing installed.""" from distilabel.models.llms import OpenAILLM from distilabel.pipeline import Pipeline from distilabel.steps import LoadDataFromDicts from distilabel.steps.tasks import TextGeneration rows = [] with open(args.seeds, encoding="utf-8") as handle: for line in handle: line = line.strip() if line: row = json.loads(line) rows.append({"id": row["id"], "instruction": row["prompt"]}) if args.limit: rows = rows[: args.limit] # Distilabel's OpenAI client reads OPENAI_API_KEY from the environment. A local # server that wants no key still wants the header to exist, so set a placeholder # rather than leaving it unset and getting an unhelpful error. os.environ.setdefault("OPENAI_API_KEY", args.api_key or "not-needed-locally") with Pipeline(name="local-teacher-and-judge") as pipeline: load = LoadDataFromDicts(data=rows, batch_size=args.batch_size) generate = TextGeneration( name="teacher", llm=OpenAILLM( model=args.teacher, base_url=args.base_url, generation_kwargs={"temperature": args.temperature, "top_p": args.top_p, "max_new_tokens": args.max_tokens}, ), input_batch_size=args.batch_size, ) judge = TextGeneration( name="judge", llm=OpenAILLM( model=args.judge, base_url=args.base_url, generation_kwargs={"temperature": 0.0, "max_new_tokens": 200}, ), template=JUDGE_TEMPLATE, columns=["instruction", "generation"], input_batch_size=args.batch_size, ) load >> generate >> judge return pipeline, len(rows) def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--seeds", default="seeds/prompts.jsonl") parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1") parser.add_argument("--api-key", default=None) parser.add_argument("--teacher", default="local/chat") parser.add_argument("--judge", default="local/chat", help="a different and preferably larger model than the teacher; " "a model grading its own answers scores them generously") parser.add_argument("--temperature", type=float, default=0.7) parser.add_argument("--top-p", type=float, default=0.8) parser.add_argument("--max-tokens", type=int, default=768) parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--limit", type=int, default=None) parser.add_argument("--out-dir", default="distilabel-out") parser.add_argument("--print-only", action="store_true", help="print the pipeline shape and exit without importing distilabel") args = parser.parse_args() if args.print_only: print("LoadDataFromDicts(data=[{id, instruction}, ...])") print(" >> TextGeneration(name='teacher', llm=OpenAILLM(model, base_url))") print(" >> TextGeneration(name='judge', llm=OpenAILLM(...), template=JUDGE_TEMPLATE,") print(" columns=['instruction', 'generation'])") print() print("Outputs of TextGeneration: 'generation' and 'model_name'.") print("The judge step reads the first step's 'generation' column through its template,") print("which is the whole reason the two steps can be chained without glue code.") print() print("Judge template:") print(JUDGE_TEMPLATE) return if args.judge == args.teacher: print("warning: the judge and the teacher are the same model. Self-preference bias " "makes the scores optimistic; Part 10's judging lesson measures this.") pipeline, count = build_pipeline(args) print(f"running the pipeline over {count} prompt(s)") distiset = pipeline.run(use_cache=False) out = Path(args.out_dir) out.mkdir(parents=True, exist_ok=True) distiset.save_to_disk(str(out)) print(f"written to {out}") print("Read a handful of rows before you train on any of them. A pipeline that runs " "without error is not the same as a pipeline that produced usable data, and " "filter-and-dedupe.py is still the step that decides which is which.") if __name__ == "__main__": main()