Spaces:
Running
Running
| """Benchmark LM Studio models against a curated AWS-RL eval set. | |
| Picks one example per (tier, source) combo from the val split, sends each | |
| prompt to every chat model loaded in LM Studio, and reports which model is | |
| the best candidate for SFT + GRPO. | |
| Usage: | |
| .venv/bin/python data/eval_lm_studio_models.py \\ | |
| --base-url http://localhost:1234/v1 \\ | |
| --val data/sft/aws_rl_sft.val.jsonl | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import re | |
| import time | |
| from pathlib import Path | |
| from typing import Any | |
| from openai import OpenAI | |
| EMBEDDING_HINT = ("embed", "embedding") | |
| def load_eval_set(val_path: Path, max_per_combo: int = 1) -> list[dict]: | |
| """One row per (tier, source) combo from the val JSONL.""" | |
| rows = [json.loads(l) for l in open(val_path)] | |
| seen: dict[tuple, int] = {} | |
| picks: list[dict] = [] | |
| for r in rows: | |
| key = (r["difficulty"], r["source"]) | |
| seen[key] = seen.get(key, 0) + 1 | |
| if seen[key] <= max_per_combo: | |
| picks.append(r) | |
| return picks | |
| def list_chat_models(client: OpenAI) -> list[str]: | |
| """Return chat-capable model ids (skip embeddings).""" | |
| out: list[str] = [] | |
| for m in client.models.list().data: | |
| mid = m.id.lower() | |
| if any(h in mid for h in EMBEDDING_HINT): | |
| continue | |
| out.append(m.id) | |
| return out | |
| def call_model( | |
| client: OpenAI, | |
| model: str, | |
| messages: list[dict], | |
| max_tokens: int = 120, | |
| timeout: float = 60.0, | |
| ) -> tuple[str, float, str | None]: | |
| """Return (completion_text, latency_s, error_or_None).""" | |
| t0 = time.time() | |
| try: | |
| resp = client.chat.completions.create( | |
| model=model, | |
| messages=messages, | |
| max_tokens=max_tokens, | |
| temperature=0.0, | |
| timeout=timeout, | |
| ) | |
| text = (resp.choices[0].message.content or "").strip() | |
| return text, time.time() - t0, None | |
| except Exception as exc: | |
| return "", time.time() - t0, f"{type(exc).__name__}: {exc}" | |
| def extract_command(raw: str) -> str: | |
| """Strip markdown fences, code blocks, leading/trailing prose to get the command.""" | |
| text = raw.strip() | |
| # Strip ```...``` fences | |
| if text.startswith("```"): | |
| lines = text.split("\n") | |
| text = "\n".join(l for l in lines if not l.startswith("```")).strip() | |
| # Take first line that starts with 'aws ' | |
| for line in text.split("\n"): | |
| line = line.strip() | |
| if line.startswith("aws "): | |
| return line | |
| return text # no aws line found — return as-is for diagnosis | |
| def score( | |
| completion: str, | |
| expected: str, | |
| ) -> dict[str, Any]: | |
| extracted = extract_command(completion) | |
| raw_stripped = completion.strip() | |
| return { | |
| "format_ok": raw_stripped.startswith("aws "), | |
| "format_ok_after_extract": extracted.startswith("aws "), | |
| "exact_match": extracted == expected.strip(), | |
| "service_match": ( | |
| extracted.split()[1:2] == expected.split()[1:2] | |
| if len(extracted.split()) >= 2 and len(expected.split()) >= 2 | |
| else False | |
| ), | |
| "operation_match": ( | |
| extracted.split()[2:3] == expected.split()[2:3] | |
| if len(extracted.split()) >= 3 and len(expected.split()) >= 3 | |
| else False | |
| ), | |
| "raw_len_chars": len(completion), | |
| "extracted": extracted[:120], | |
| } | |
| def run_benchmark( | |
| client: OpenAI, | |
| models: list[str], | |
| eval_set: list[dict], | |
| ) -> dict[str, list[dict]]: | |
| results: dict[str, list[dict]] = {m: [] for m in models} | |
| for i, task in enumerate(eval_set): | |
| expected = task["messages"][2]["content"] | |
| messages_in = task["messages"][:2] # system + user, no assistant | |
| tier = task["difficulty"] | |
| source = task["source"] | |
| print(f"\n[{i+1}/{len(eval_set)}] tier={tier} source={source} task_id={task['task_id']}") | |
| print(f" expected: {expected[:90]!r}") | |
| for model in models: | |
| completion, latency, err = call_model(client, model, messages_in) | |
| if err: | |
| row = { | |
| "tier": tier, "source": source, "task_id": task["task_id"], | |
| "completion": "", "error": err, "latency_s": round(latency, 2), | |
| "format_ok": False, "format_ok_after_extract": False, | |
| "exact_match": False, "service_match": False, | |
| "operation_match": False, "raw_len_chars": 0, "extracted": "", | |
| } | |
| else: | |
| s = score(completion, expected) | |
| row = { | |
| "tier": tier, "source": source, "task_id": task["task_id"], | |
| "completion": completion, "error": None, | |
| "latency_s": round(latency, 2), | |
| **s, | |
| } | |
| results[model].append(row) | |
| flag = "✓" if row.get("exact_match") else ("~" if row.get("format_ok_after_extract") else "✗") | |
| print(f" {flag} {model:<35} {latency:5.1f}s {row.get('extracted','')[:70]!r}") | |
| return results | |
| def aggregate(results: dict[str, list[dict]]) -> list[dict]: | |
| agg = [] | |
| for model, rows in results.items(): | |
| n = len(rows) | |
| if n == 0: | |
| continue | |
| agg.append({ | |
| "model": model, | |
| "n": n, | |
| "errors": sum(1 for r in rows if r.get("error")), | |
| "format_ok_pct": round(sum(1 for r in rows if r["format_ok"]) / n, 2), | |
| "format_after_extract_pct": round( | |
| sum(1 for r in rows if r["format_ok_after_extract"]) / n, 2 | |
| ), | |
| "exact_match_pct": round(sum(1 for r in rows if r["exact_match"]) / n, 2), | |
| "service_match_pct": round(sum(1 for r in rows if r["service_match"]) / n, 2), | |
| "operation_match_pct": round(sum(1 for r in rows if r["operation_match"]) / n, 2), | |
| "avg_latency_s": round(sum(r["latency_s"] for r in rows) / n, 2), | |
| "avg_len_chars": round(sum(r["raw_len_chars"] for r in rows) / n, 1), | |
| }) | |
| agg.sort( | |
| key=lambda d: (d["format_after_extract_pct"], d["exact_match_pct"], -d["avg_latency_s"]), | |
| reverse=True, | |
| ) | |
| return agg | |
| def print_table(agg: list[dict]) -> None: | |
| print("\n" + "=" * 110) | |
| print(f"{'Model':<36} {'n':>3} {'errs':>4} {'fmt%':>5} {'+xtr%':>6} {'exact%':>7} {'svc%':>5} {'op%':>5} {'lat':>5} {'len':>5}") | |
| print("-" * 110) | |
| for r in agg: | |
| print( | |
| f"{r['model']:<36} {r['n']:>3} {r['errors']:>4} " | |
| f"{int(r['format_ok_pct']*100):>4}% {int(r['format_after_extract_pct']*100):>5}% " | |
| f"{int(r['exact_match_pct']*100):>6}% {int(r['service_match_pct']*100):>4}% " | |
| f"{int(r['operation_match_pct']*100):>4}% {r['avg_latency_s']:>4.1f}s {int(r['avg_len_chars']):>4}" | |
| ) | |
| print("=" * 110) | |
| print("Column legend:") | |
| print(" fmt% — raw output starts with 'aws ' (no preamble, no fences)") | |
| print(" +xtr% — starts with 'aws ' after stripping fences/prose") | |
| print(" exact% — extracted command matches canonical exactly") | |
| print(" svc% — same AWS service (e.g. s3, dynamodb)") | |
| print(" op% — same operation (e.g. create-bucket)") | |
| print(" lat — mean seconds per call | len — mean raw chars") | |
| def main() -> None: | |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) | |
| ap.add_argument("--base-url", default="http://localhost:1234/v1") | |
| ap.add_argument("--val", type=Path, default=Path("data/sft/aws_rl_sft.val.jsonl")) | |
| ap.add_argument("--out", type=Path, default=Path("data/sft/model_eval_results.json")) | |
| ap.add_argument("--max-per-combo", type=int, default=1) | |
| args = ap.parse_args() | |
| client = OpenAI(base_url=args.base_url, api_key="lm-studio") | |
| models = list_chat_models(client) | |
| print(f"Found {len(models)} chat models: {models}") | |
| eval_set = load_eval_set(args.val, args.max_per_combo) | |
| print(f"Eval set: {len(eval_set)} prompts (one per (tier, source) combo)") | |
| results = run_benchmark(client, models, eval_set) | |
| agg = aggregate(results) | |
| print_table(agg) | |
| with open(args.out, "w") as f: | |
| json.dump({"aggregate": agg, "per_call": results}, f, indent=2) | |
| print(f"\nFull results saved to {args.out}") | |
| if __name__ == "__main__": | |
| main() | |