Spaces:
Sleeping
Sleeping
| """Task 2: turn a natural-language training-run request into a valid oumi YAML. | |
| The agent gets two tools: get_schema (discover the supported config subset) and | |
| validate_config (YAML parse + JSON Schema check). Dealt requests are rendered | |
| from a sampled field spec, so ground truth is the spec itself: scoring checks | |
| schema validity plus exact per-field match on the targeted values. | |
| """ | |
| import json | |
| import random | |
| import re | |
| import jsonschema | |
| import yaml | |
| from app.config import DATA_DIR | |
| from app.tasks.base import Sample, Task | |
| SCHEMA = json.loads((DATA_DIR / "oumi_training_schema.json").read_text()) | |
| _VALIDATOR = jsonschema.Draft202012Validator(SCHEMA) | |
| SYSTEM_PROMPT = """\ | |
| You are the Oumi Config Copilot. You turn a user's plain-English description of a \ | |
| fine-tuning run into a valid oumi training config (YAML). | |
| Rules: | |
| - Only use fields that exist in the config schema. Call get_schema to see the \ | |
| schema for a section before using unfamiliar fields. | |
| - Supervised fine-tuning uses training.trainer_type: TRL_SFT. | |
| - LoRA/QLoRA runs set training.use_peft: true and configure the peft section. \ | |
| QLoRA additionally sets peft.q_lora: true. | |
| - For a local JSONL file of chat conversations, use dataset_name: "text_sft_jsonl" \ | |
| with dataset_path pointing at the file. | |
| - Include the fields the user asked for plus anything required by the schema. \ | |
| Don't pad the config with unrelated defaults. | |
| - Before giving your final answer, call validate_config on your draft and fix any \ | |
| errors it reports. | |
| - Your final message must be exactly one fenced ```yaml code block containing the \ | |
| config, with no text after it.\ | |
| """ | |
| TOOLS = [ | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "get_schema", | |
| "description": "Return the JSON Schema for a section of the oumi training config (or the whole thing).", | |
| "parameters": { | |
| "type": "object", | |
| "properties": { | |
| "section": { | |
| "type": "string", | |
| "enum": ["all", "model", "data", "training", "peft"], | |
| } | |
| }, | |
| "required": ["section"], | |
| }, | |
| }, | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "validate_config", | |
| "description": "Validate a candidate oumi training config. Returns {valid, errors}.", | |
| "parameters": { | |
| "type": "object", | |
| "properties": { | |
| "yaml_config": {"type": "string", "description": "The full YAML config text"} | |
| }, | |
| "required": ["yaml_config"], | |
| }, | |
| }, | |
| }, | |
| ] | |
| def validate(config: object) -> list[str]: | |
| errors = sorted(_VALIDATOR.iter_errors(config), key=lambda e: list(e.absolute_path)) | |
| return [ | |
| f"at {'.'.join(str(p) for p in e.absolute_path) or '<root>'}: {e.message}" | |
| for e in errors | |
| ][:12] | |
| def execute_tool(name: str, args: dict) -> str: | |
| if name == "get_schema": | |
| section = args.get("section", "all") | |
| sub = SCHEMA if section == "all" else SCHEMA["properties"].get(section, {}) | |
| return json.dumps(sub, indent=1) | |
| if name == "validate_config": | |
| try: | |
| config = yaml.safe_load(args.get("yaml_config", "")) | |
| except yaml.YAMLError as e: | |
| return json.dumps({"valid": False, "errors": [f"YAML parse error: {e}"]}) | |
| if not isinstance(config, dict): | |
| return json.dumps({"valid": False, "errors": ["config must be a YAML mapping"]}) | |
| errors = validate(config) | |
| return json.dumps({"valid": not errors, "errors": errors}) | |
| return json.dumps({"error": f"unknown tool {name}"}) | |
| # --- Request generation ------------------------------------------------------ | |
| MODELS = [ | |
| "Qwen/Qwen3-4B-Instruct-2507", "Qwen/Qwen2.5-1.5B-Instruct", | |
| "meta-llama/Llama-3.2-3B-Instruct", "meta-llama/Llama-3.1-8B-Instruct", | |
| "HuggingFaceTB/SmolLM2-1.7B-Instruct", "google/gemma-3-4b-it", | |
| "microsoft/Phi-4-mini-instruct", "mistralai/Mistral-7B-Instruct-v0.3", | |
| ] | |
| HUB_DATASETS = [ | |
| "yahma/alpaca-cleaned", "HuggingFaceH4/ultrachat_200k", | |
| "databricks/databricks-dolly-15k", "OpenAssistant/oasst1", | |
| ] | |
| LOCAL_PATHS = [ | |
| "data/train.jsonl", "/mnt/data/conversations.jsonl", "datasets/support_chats.jsonl", | |
| "data/curated/sft_round2.jsonl", "exports/traces_clean.jsonl", | |
| ] | |
| PROJECT_SLUGS = ["support-bot", "triage-v2", "summarizer", "qa-assistant", "intent-router"] | |
| OPENERS = [ | |
| "", "Hey! ", "Quick one: ", "Setting up a new experiment. ", | |
| "Time for another ablation. ", "For the {slug} project: ", | |
| "Need a config for tomorrow's run. ", | |
| ] | |
| def generate(seed: int) -> tuple[dict, str]: | |
| """Deterministically generate (target field spec, natural-language request).""" | |
| rng = random.Random(seed) | |
| fields: dict[str, object] = {} | |
| phrases: list[str] = [] | |
| fields["training.trainer_type"] = "TRL_SFT" | |
| model = rng.choice(MODELS) | |
| fields["model.model_name"] = model | |
| kind = rng.choice(["full", "lora", "lora", "qlora"]) # lora-heavy mix | |
| kind_phrase = { | |
| "full": rng.choice([f"do a full-parameter SFT run on {model}", | |
| f"fully fine-tune {model} (no adapters)"]), | |
| "lora": rng.choice([f"fine-tune {model} with LoRA", | |
| f"do a LoRA SFT run on {model}"]), | |
| "qlora": rng.choice([f"fine-tune {model} with QLoRA (4-bit)", | |
| f"do a 4-bit QLoRA run on {model}"]), | |
| }[kind] | |
| fields["training.use_peft"] = kind != "full" | |
| if kind == "qlora": | |
| fields["peft.q_lora"] = True | |
| fields["peft.q_lora_bits"] = 4 | |
| if kind in ("lora", "qlora"): | |
| if rng.random() < 0.7: | |
| r = rng.choice([8, 16, 32, 64]) | |
| fields["peft.lora_r"] = r | |
| phrases.append(rng.choice([f"set the LoRA rank to {r}", f"use lora_r {r}"])) | |
| if rng.random() < 0.5: | |
| alpha = rng.choice([r, 2 * r]) | |
| fields["peft.lora_alpha"] = alpha | |
| phrases.append(f"with alpha {alpha}") | |
| if rng.random() < 0.6: | |
| path = rng.choice(LOCAL_PATHS) | |
| fields["data.train.datasets[0].dataset_name"] = "text_sft_jsonl" | |
| fields["data.train.datasets[0].dataset_path"] = path | |
| data_phrase = rng.choice([ | |
| f"on my conversation data at {path} (JSONL, oumi chat format)", | |
| f"using the local JSONL file {path} as training data", | |
| ]) | |
| if rng.random() < 0.25: | |
| val = re.sub(r"\.jsonl$", "_val.jsonl", path) | |
| fields["data.validation.datasets[0].dataset_name"] = "text_sft_jsonl" | |
| fields["data.validation.datasets[0].dataset_path"] = val | |
| phrases.append(f"use {val} as the validation set") | |
| else: | |
| ds = rng.choice(HUB_DATASETS) | |
| fields["data.train.datasets[0].dataset_name"] = ds | |
| data_phrase = rng.choice([f"on the {ds} dataset", f"using {ds} from the Hub"]) | |
| if rng.random() < 0.2: | |
| n = rng.choice([1000, 5000, 10000, 20000]) | |
| fields["data.train.datasets[0].sample_count"] = n | |
| phrases.append(f"cap the training data at {n} examples") | |
| if rng.random() < 0.2: | |
| fields["data.train.pack"] = True | |
| phrases.append("pack sequences") | |
| if rng.random() < 0.65: | |
| n = rng.choice([1, 2, 3, 4, 5]) | |
| fields["training.num_train_epochs"] = n | |
| phrases.append(rng.choice([f"train for {n} epoch{'s' if n > 1 else ''}", | |
| f"{n} epoch{'s' if n > 1 else ''}"])) | |
| else: | |
| n = rng.choice([100, 200, 500, 1000, 2000, 3000]) | |
| fields["training.max_steps"] = n | |
| phrases.append(f"train for {n} steps") | |
| lr = rng.choice(["1e-5", "2e-5", "5e-5", "1e-4", "2e-4", "3e-4", "5e-4"]) | |
| fields["training.learning_rate"] = float(lr) | |
| phrases.append(rng.choice([f"learning rate {lr}", f"lr {lr}"])) | |
| for prob, path, choices, phrase_fn in [ | |
| (0.6, "training.per_device_train_batch_size", [1, 2, 4, 8, 16], | |
| lambda v: rng.choice([f"batch size {v} per device", f"per-device batch size of {v}"])), | |
| (0.4, "training.gradient_accumulation_steps", [2, 4, 8, 16, 32], | |
| lambda v: f"gradient accumulation {v}"), | |
| (0.5, "training.mixed_precision_dtype", ["bf16", "fp16"], | |
| lambda v: f"in {v} mixed precision"), | |
| (0.4, "training.output_dir", [f"output/{s}" for s in PROJECT_SLUGS], | |
| lambda v: rng.choice([f"write everything to {v}", f"output dir {v}"])), | |
| (0.3, "training.seed", [0, 1, 7, 42, 123, 1234], | |
| lambda v: f"seed {v}"), | |
| (0.35, "training.lr_scheduler_type", ["cosine", "linear", "constant"], | |
| lambda v: f"{v} LR schedule"), | |
| (0.25, "training.warmup_steps", [10, 50, 100, 200], | |
| lambda v: f"{v} warmup steps"), | |
| (0.2, "training.weight_decay", [0.01, 0.1], | |
| lambda v: f"weight decay {v}"), | |
| (0.25, "training.save_steps", [100, 250, 500, 1000], | |
| lambda v: f"checkpoint every {v} steps"), | |
| (0.2, "training.logging_steps", [10, 25, 50, 100], | |
| lambda v: f"log every {v} steps"), | |
| (0.25, "training.enable_gradient_checkpointing", [True], | |
| lambda v: "turn on gradient checkpointing"), | |
| (0.3, "model.model_max_length", [1024, 2048, 4096, 8192], | |
| lambda v: rng.choice([f"max sequence length {v}", f"context length {v}"])), | |
| (0.2, "training.run_name", [f"{s}-{n}" for s in PROJECT_SLUGS for n in ("v1", "v2", "exp3")], | |
| lambda v: f"call the run {v}"), | |
| ]: | |
| if rng.random() < prob: | |
| value = rng.choice(choices) | |
| fields[path] = value | |
| phrases.append(phrase_fn(value)) | |
| if rng.random() < 0.4: | |
| n = rng.choice([50, 100, 200, 250, 500]) | |
| fields["training.eval_strategy"] = "steps" | |
| fields["training.eval_steps"] = n | |
| phrases.append(f"evaluate every {n} steps") | |
| rng.shuffle(phrases) | |
| opener = rng.choice(OPENERS).format(slug=rng.choice(PROJECT_SLUGS)) | |
| text = f"{opener}I want to {kind_phrase} {data_phrase}." | |
| if phrases: | |
| text += " " + ", ".join(phrases) + "." | |
| return fields, text | |
| def get_path(config: dict, path: str) -> object: | |
| """Navigate 'a.b[0].c'-style paths; returns None when absent.""" | |
| current: object = config | |
| for part in re.findall(r"[^.\[\]]+|\[\d+\]", path): | |
| if part.startswith("["): | |
| if not isinstance(current, list): | |
| return None | |
| idx = int(part[1:-1]) | |
| if idx >= len(current): | |
| return None | |
| current = current[idx] | |
| else: | |
| if not isinstance(current, dict): | |
| return None | |
| current = current.get(part) | |
| if current is None: | |
| return None | |
| return current | |
| def values_match(expected: object, got: object) -> bool: | |
| if isinstance(expected, bool) or isinstance(got, bool): | |
| return isinstance(expected, bool) and isinstance(got, bool) and expected is got | |
| if isinstance(expected, (int, float)) and isinstance(got, (int, float)): | |
| return float(expected) == float(got) | |
| if isinstance(expected, str) and isinstance(got, str): | |
| return expected.strip().lower() == got.strip().lower() | |
| return expected == got | |
| # --- Task hooks --------------------------------------------------------------- | |
| _rng = random.Random() | |
| def sample() -> Sample: | |
| seed = _rng.randrange(1, 10**9) | |
| _, text = generate(seed) | |
| return Sample(input_id=f"req-{seed}", text=text) | |
| def lookup_truth(input_id: str) -> dict | None: | |
| m = re.fullmatch(r"req-(\d+)", input_id or "") | |
| if not m: | |
| return None | |
| fields, _ = generate(int(m.group(1))) | |
| return {"fields": fields} | |
| def parse_output(text: str) -> dict: | |
| blocks = re.findall(r"```(?:ya?ml)?\s*\n(.*?)```", text, re.DOTALL) | |
| yaml_text = blocks[-1].strip() if blocks else text.strip() | |
| config, parse_error = None, None | |
| try: | |
| loaded = yaml.safe_load(yaml_text) | |
| if isinstance(loaded, dict): | |
| config = loaded | |
| else: | |
| parse_error = "config is not a YAML mapping" | |
| except yaml.YAMLError as e: | |
| parse_error = f"YAML parse error: {e}" | |
| errors = [parse_error] if parse_error else validate(config) | |
| return {"yaml_text": yaml_text, "config": config, "errors": errors} | |
| def score(truth: dict, parsed: dict) -> dict[str, float]: | |
| config = parsed["config"] | |
| scores = { | |
| "yaml_valid": float(config is not None), | |
| "schema_valid": float(config is not None and not parsed["errors"]), | |
| } | |
| fields = truth["fields"] | |
| matched = sum( | |
| values_match(expected, get_path(config or {}, path)) | |
| for path, expected in fields.items() | |
| ) | |
| scores["field_match"] = matched / len(fields) | |
| return scores | |
| def present(parsed: dict, truth: dict | None) -> dict: | |
| out = { | |
| "yaml": parsed["yaml_text"], | |
| "yaml_valid": parsed["config"] is not None, | |
| "schema_valid": parsed["config"] is not None and not parsed["errors"], | |
| "errors": parsed["errors"], | |
| } | |
| if truth: | |
| config = parsed["config"] or {} | |
| out["fields"] = [ | |
| { | |
| "path": path, | |
| "expected": expected, | |
| "got": get_path(config, path), | |
| "match": values_match(expected, get_path(config, path)), | |
| } | |
| for path, expected in truth["fields"].items() | |
| ] | |
| return out | |
| TASK = Task( | |
| order=2, | |
| id="config-copilot", | |
| title="Config Copilot", | |
| tagline="Describe a training run, get a valid oumi config", | |
| system_prompt=SYSTEM_PROMPT, | |
| ui={ | |
| "output": "config", | |
| "input_label": "Training run request", | |
| "placeholder": "Deal a request, or describe the fine-tuning run you want...", | |
| "deal_label": "Deal me a request", | |
| }, | |
| sample=sample, | |
| lookup_truth=lookup_truth, | |
| parse_output=parse_output, | |
| score=score, | |
| present=present, | |
| tools=TOOLS, | |
| execute_tool=execute_tool, | |
| # Routed to the "oumi" backend (an Oumi proxy deployment) when LLM_OUMI_* is | |
| # configured; falls back to the default backend otherwise. The other tasks | |
| # stay on the default backend. See config.backend() and .env.example. | |
| backend="oumi", | |
| ) | |