159 lines
7.0 KiB
Python
159 lines
7.0 KiB
Python
"""Offline DSPy optimization of the question catalog (dev dependency: dspy).
|
|
|
|
python -m tools.answers.catalog_optimize [--threads 2] [--provisional]
|
|
|
|
Reads the labeled questions (catalog_examples.json), splits them into train and
|
|
dev sets deterministically, compiles a few-shot program with DSPy
|
|
(BootstrapFewShotWithRandomSearch) against the configured model and writes
|
|
question_catalog.json: the instructions, the selected demonstrations and the
|
|
dev scores of the RUNTIME classifier (the same code path the backend uses).
|
|
--provisional writes a labeled few-shot program without any model call.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from tools.answers.context_policy import normalized
|
|
from tools.answers.question_catalog import (CatalogDecision, INSTRUCTIONS, EXAMPLES_FILE, PROGRAM_FILE,
|
|
TYPE_DESCRIPTIONS, checked, classify, _cache)
|
|
|
|
|
|
def load_examples():
|
|
rows = json.loads(EXAMPLES_FILE.read_text(encoding="utf-8"))
|
|
examples = []
|
|
for row in rows:
|
|
fields = {k: v for k, v in row.items() if k != "q"}
|
|
examples.append((row["q"], CatalogDecision(**fields)))
|
|
dev = [e for i, e in enumerate(examples) if i % 4 == 3]
|
|
train = [e for i, e in enumerate(examples) if i % 4 != 3]
|
|
return train, dev
|
|
|
|
|
|
def _same_phrase(a, b):
|
|
if not a and not b:
|
|
return True
|
|
if not a or not b:
|
|
return False
|
|
a, b = normalized(a).strip(), normalized(b).strip()
|
|
return a in b or b in a
|
|
|
|
|
|
def score(gold: CatalogDecision, got: CatalogDecision | None) -> float:
|
|
if got is None:
|
|
return 0.0
|
|
value = 0.5 * (gold.type == got.type)
|
|
value += 0.1 * (len(gold.people) == len(got.people))
|
|
value += 0.15 * _same_phrase(gold.topic, got.topic)
|
|
value += 0.05 * _same_phrase(gold.exclude, got.exclude)
|
|
value += 0.05 * (gold.action == got.action)
|
|
value += 0.05 * (set(gold.facets) == set(got.facets))
|
|
value += 0.05 * (gold.presentation == got.presentation)
|
|
value += 0.05 * (set(gold.work_types) == set(got.work_types))
|
|
return round(value, 4)
|
|
|
|
|
|
def decision_json(decision: CatalogDecision) -> dict:
|
|
return decision.model_dump(exclude_defaults=True)
|
|
|
|
|
|
def provisional(train):
|
|
"""Two labeled demonstrations per type (DSPy LabeledFewShot equivalent)."""
|
|
demos, per_type = [], {}
|
|
for question, decision in train:
|
|
if per_type.get(decision.type, 0) < 2:
|
|
per_type[decision.type] = per_type.get(decision.type, 0) + 1
|
|
demos.append({"question": question, "decision": decision_json(decision)})
|
|
return {"instructions": INSTRUCTIONS, "demos": demos, "optimizer": "labeled-provisional"}
|
|
|
|
|
|
def optimize(train, dev, threads):
|
|
import dspy
|
|
from tools.answers.model_client import model_configuration
|
|
config = model_configuration()
|
|
api_base = config["endpoint"].rsplit("/chat/completions", 1)[0]
|
|
lm = dspy.LM("openai/" + config["model"], api_base=api_base, api_key=config["key"], temperature=0,
|
|
max_tokens=400, cache=True, num_retries=2, model_type="chat")
|
|
dspy.configure(lm=lm, adapter=dspy.JSONAdapter())
|
|
|
|
class Classify(dspy.Signature):
|
|
"""Classify a ZP Wiki question."""
|
|
question: str = dspy.InputField()
|
|
decision: CatalogDecision = dspy.OutputField()
|
|
|
|
signature = Classify.with_instructions(INSTRUCTIONS)
|
|
trainset = [dspy.Example(question=q, decision=d).with_inputs("question") for q, d in train]
|
|
devset = [dspy.Example(question=q, decision=d).with_inputs("question") for q, d in dev]
|
|
|
|
def metric(example, prediction, trace=None):
|
|
got = getattr(prediction, "decision", None)
|
|
if isinstance(got, dict):
|
|
got = CatalogDecision.model_validate(got)
|
|
if got is not None:
|
|
got = checked(got, example.question)
|
|
value = score(example.decision, got)
|
|
return value >= 0.9 if trace is not None else value
|
|
|
|
from dspy.teleprompt import BootstrapFewShotWithRandomSearch
|
|
optimizer = BootstrapFewShotWithRandomSearch(metric=metric, max_bootstrapped_demos=4, max_labeled_demos=20,
|
|
num_candidate_programs=4, num_threads=threads)
|
|
compiled = optimizer.compile(dspy.Predict(signature), trainset=trainset, valset=devset)
|
|
demos = []
|
|
for demo in compiled.demos:
|
|
decision = demo.decision
|
|
if isinstance(decision, dict):
|
|
decision = CatalogDecision.model_validate(decision)
|
|
demos.append({"question": demo.question, "decision": decision_json(decision)})
|
|
covered = {d["decision"]["type"] for d in demos}
|
|
# Every type keeps at least one labeled example; the optimizer may drop rare ones.
|
|
for question, decision in train:
|
|
if decision.type not in covered:
|
|
demos.append({"question": question, "decision": decision_json(decision)})
|
|
covered.add(decision.type)
|
|
return {"instructions": compiled.signature.instructions, "demos": demos,
|
|
"optimizer": "dspy.BootstrapFewShotWithRandomSearch",
|
|
"dspy_version": dspy.__version__}
|
|
|
|
|
|
def evaluate(program, dev):
|
|
"""Scores of the runtime classifier (not DSPy) on the dev questions."""
|
|
_cache.clear()
|
|
results = []
|
|
for question, gold in dev:
|
|
got = classify(question, program=program)
|
|
results.append({"question": question, "gold": gold.type, "got": got.type if got else None,
|
|
"score": score(gold, got)})
|
|
accuracy = sum(r["gold"] == r["got"] for r in results) / len(results)
|
|
return {"type_accuracy": round(accuracy, 4),
|
|
"mean_score": round(sum(r["score"] for r in results) / len(results), 4),
|
|
"errors": [r for r in results if r["gold"] != r["got"]]}
|
|
|
|
|
|
def main():
|
|
from dotenv import load_dotenv
|
|
load_dotenv(Path(__file__).resolve().parents[2] / ".env", override=False)
|
|
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
|
|
parser.add_argument("--threads", type=int, default=2)
|
|
parser.add_argument("--provisional", action="store_true")
|
|
parser.add_argument("--skip-eval", action="store_true")
|
|
args = parser.parse_args()
|
|
train, dev = load_examples()
|
|
started = time.time()
|
|
program = provisional(train) if args.provisional else optimize(train, dev, args.threads)
|
|
program.update({"types": list(TYPE_DESCRIPTIONS), "train_size": len(train), "dev_size": len(dev),
|
|
"built_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())})
|
|
if not args.skip_eval:
|
|
program["dev_evaluation"] = evaluate(program, dev)
|
|
program["seconds"] = round(time.time() - started)
|
|
PROGRAM_FILE.write_text(json.dumps(program, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
|
|
print(json.dumps({k: program[k] for k in ("optimizer", "train_size", "dev_size", "seconds")}
|
|
| {"demos": len(program["demos"]), **{k: v for k, v in program.get("dev_evaluation", {}).items()
|
|
if k != "errors"}}, ensure_ascii=False))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|