contimp-app / app /tasks /pr_area.py
lefft's picture
contimp-app v0.1.0
910dadd verified
Raw
History Blame Contribute Delete
3.34 kB
"""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,
)