Spaces:
Sleeping
Sleeping
| """Task 1: classify which oumi subsystem a PR touches, from its title + body.""" | |
| import json | |
| import random | |
| import re | |
| from app.config import DATA_DIR | |
| from app.tasks.base import Sample, Task | |
| AREAS = [ | |
| "training", "inference", "data", "evaluation", "configs", | |
| "launcher", "cli", "docs", "infra", "other", | |
| ] | |
| SYSTEM_PROMPT = """\ | |
| You are a pull-request triage classifier for oumi, an open-source library for \ | |
| training, evaluating, and deploying foundation models. | |
| Given a PR's title and body, classify which subsystem of the repo it primarily \ | |
| changes. The labels: | |
| - training: training loop, trainers, tuners, callbacks, collators, distributed/performance code | |
| - inference: inference engines, serving, deployment code | |
| - data: dataset classes, tokenizers, processors, data synthesis, conversation types | |
| - evaluation: evaluation harnesses, judges, analyzers, metrics | |
| - configs: YAML recipes for models/jobs, or the config schema code (params classes) | |
| - launcher: launching/managing jobs on clouds and clusters (SkyPilot, Slurm, ...) | |
| - cli: the `oumi` command-line interface | |
| - docs: documentation, notebooks, markdown files | |
| - infra: CI/CD, GitHub workflows, dependency bumps, build tooling, repo maintenance | |
| - other: cross-cutting utilities, model architecture code, or anything not above | |
| Respond with exactly one JSON object on a single line, no other text: | |
| {"area": "<label>"}\ | |
| """ | |
| _RECORDS: dict[str, dict] = {} | |
| with (DATA_DIR / "prs.jsonl").open() as _f: | |
| for _line in _f: | |
| _r = json.loads(_line) | |
| _RECORDS[f"pr-{_r['number']}"] = _r | |
| _rng = random.Random() | |
| def render_pr(record: dict) -> str: | |
| body = (record.get("body") or "").strip() | |
| text = f"Title: {record['title']}" | |
| if body: | |
| text += f"\n\nBody:\n{body}" | |
| return text | |
| def sample() -> Sample: | |
| input_id = _rng.choice(list(_RECORDS)) | |
| return Sample(input_id=input_id, text=render_pr(_RECORDS[input_id])) | |
| def lookup_truth(input_id: str) -> dict | None: | |
| record = _RECORDS.get(input_id) | |
| if record is None: | |
| return None | |
| return {"area": record["area"], "pr_number": record["number"]} | |
| def parse_output(text: str) -> dict: | |
| m = re.search(r"\{[^{}]*\}", text, re.DOTALL) | |
| area = None | |
| if m: | |
| try: | |
| value = json.loads(m.group()).get("area") | |
| if isinstance(value, str) and value.strip().lower() in AREAS: | |
| area = value.strip().lower() | |
| except json.JSONDecodeError: | |
| pass | |
| return {"area": area, "raw": text} | |
| def score(truth: dict, parsed: dict) -> dict[str, float]: | |
| return { | |
| "format_ok": float(parsed["area"] is not None), | |
| "exact_match": float(parsed["area"] == truth["area"]), | |
| } | |
| def present(parsed: dict, truth: dict | None) -> dict: | |
| return {"area": parsed["area"], "raw": parsed["raw"]} | |
| TASK = Task( | |
| order=1, | |
| id="pr-area", | |
| title="PR Triage", | |
| tagline="Which oumi subsystem does this PR touch?", | |
| system_prompt=SYSTEM_PROMPT, | |
| ui={ | |
| "output": "label", | |
| "labels": AREAS, | |
| "input_label": "Pull request (title + body)", | |
| "placeholder": "Deal a real oumi PR, or write your own title + body...", | |
| "deal_label": "Deal me a PR", | |
| }, | |
| sample=sample, | |
| lookup_truth=lookup_truth, | |
| parse_output=parse_output, | |
| score=score, | |
| present=present, | |
| ) | |