Adds benchmarking scripts for llm.

Author: ppak10

README.md

@@ -46,3 +46,14 @@ am workspace create test
 
 #### Example
 An example implementation can be found [here](https://github.com/ppak10/additive-manufacturing-agent)
+
+## Benchmark
+To run the benchmark commands, install the optional `benchmark` dependencies:
+
+```bash
+uv pip install "additive-manufacturing[benchmark]"
+```
+
+```bash
+am benchmark <model>
+```

pyproject.toml

@@ -33,6 +33,13 @@ Issues = "https://github.com/ppak10/additive-manufacturing/issues"
 am = "am.cli:app"
 
 [project.optional-dependencies]
+benchmark = [
+    "accelerate>=1.12.0",
+    "datasets>=4.6.0",
+    "sentence-transformers>=5.2.3",
+    "torch>=2.10.0",
+    "transformers>=5.2.0",
+]
 cuda12 = [
     "jax[cuda12]>=0.6.2",
 ]

src/am/benchmark/__init__.py

@@ -0,0 +1,58 @@
+import typer
+
+
+def register_benchmark(app: typer.Typer):
+    from pathlib import Path
+    from rich import print as rprint
+
+    from am.benchmark.llm import benchmark_llm, TASKS
+    from am.cli.options import NumProc
+    from wa.cli.options import WorkspaceOption
+
+    @app.command(name="benchmark", rich_help_panel="Benchmark Commands")
+    def benchmark(
+        model: str,
+        tasks: list[str] = typer.Option(TASKS, help="Tasks to benchmark"),
+        batch_size: int = typer.Option(
+            8, help="Number of prompts to process in parallel."
+        ),
+        url: str = typer.Option(
+            None, help="URL of a running vLLM server (e.g. http://localhost:8000/v1)."
+        ),
+        workspace_option: WorkspaceOption = None,
+        num_proc: NumProc = 1,
+    ) -> None:
+        """
+        Runs additive manufacturing benchmark for large language models.
+        """
+
+        from wa.cli.utils import get_workspace
+
+        workspace = get_workspace(workspace_option)
+
+        try:
+            workspace_folder_path = None
+            if workspace:
+                model_name = model.replace("/", "--")
+                workspace_folder = workspace.create_folder(
+                    name_or_path=Path("benchmarks") / model_name, append_timestamp=True
+                )
+                workspace_folder_path = workspace_folder.path
+
+            benchmark_llm(
+                model=model,
+                tasks=tasks,
+                batch_size=batch_size,
+                url=url,
+                num_proc=num_proc,
+                out_path=workspace_folder_path,
+            )
+
+        except Exception as e:
+            import traceback
+
+            traceback.print_exc()
+            rprint(f"⚠️ [yellow]Unable to run benchmark for {model}: {e}[/yellow]")
+            raise typer.Exit(code=1)
+
+    return benchmark

src/am/benchmark/cli.py

@@ -0,0 +1,69 @@
+import importlib
+from pathlib import Path
+
+from .constants import TASKS
+from .inference import _run_openai_compatible, _run_transformers
+from .general_knowledge_multiple_choice import (
+    _benchmark_general_knowledge_multiple_choice,
+)
+from .general_knowledge_short_answer import _benchmark_general_knowledge_short_answer
+from .melt_pool_geometry_prediction import _benchmark_melt_pool_geometry_prediction
+
+
+def benchmark_llm(
+    model: str,
+    tasks: list[str] = TASKS,
+    num_proc: int = 1,
+    batch_size: int = 8,
+    url: str | None = None,
+    out_path: Path | None = None,
+) -> dict:
+
+    missing = [
+        pkg
+        for pkg in ("datasets", "sentence_transformers", "transformers")
+        if importlib.util.find_spec(pkg) is None
+    ]
+    if missing:
+        raise ImportError(
+            f"Missing benchmark packages: {', '.join(missing)}. "
+            'Run: uv pip install "additive-manufacturing[benchmark]"'
+        )
+
+    if url is not None:
+        print(f"Using vLLM server: {url}")
+        runner = lambda questions, max_tokens: _run_openai_compatible(
+            url, model, questions, max_tokens
+        )
+    else:
+        from transformers import pipeline
+
+        print(f"Loading model: {model}")
+        llm_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            return_full_text=False,
+            device_map="auto",
+        )
+        runner = lambda questions, max_tokens: _run_transformers(
+            llm_pipeline, questions, batch_size, max_tokens
+        )
+
+    reports = {}
+    for task in tasks:
+        if task == "general_knowledge_multiple_choice":
+            reports[task] = _benchmark_general_knowledge_multiple_choice(
+                runner, model, num_proc, out_path
+            )
+        elif task == "general_knowledge_short_answer":
+            reports[task] = _benchmark_general_knowledge_short_answer(
+                runner, model, num_proc, out_path
+            )
+        elif task == "melt_pool_geometry_prediction":
+            reports[task] = _benchmark_melt_pool_geometry_prediction(
+                runner, model, num_proc, out_path
+            )
+        else:
+            print(f"Unknown task '{task}', skipping.")
+
+    return reports

src/am/benchmark/llm/__init__.py

@@ -0,0 +1,11 @@
+DATASET_NAME = "ppak10/Additive-Manufacturing-Benchmark"
+EVALUATOR_MODEL = "all-MiniLM-L6-v2"
+RUBRIC_EVALUATOR_MODEL = "cross-encoder/nli-deberta-v3-large"
+MAX_NEW_TOKENS_SHORT_ANSWER = 4096
+MAX_NEW_TOKENS_MELT_POOL = 1024
+MAX_NEW_TOKENS_MULTIPLE_CHOICE = 256
+TASKS = [
+    "general_knowledge_multiple_choice",
+    "general_knowledge_short_answer",
+    "melt_pool_geometry_prediction",
+]

src/am/benchmark/llm/constants.py

@@ -0,0 +1,123 @@
+import json
+import re
+from datetime import datetime, timezone
+from pathlib import Path
+
+from datasets import load_dataset
+
+from .constants import DATASET_NAME, MAX_NEW_TOKENS_MULTIPLE_CHOICE
+
+
+def _build_mc_prompt(row: dict) -> str:
+    lines = [
+        "You are an expert in additive manufacturing.",
+        "Answer the following multiple-choice question by responding with only the letter of the correct answer (A, B, C, or D).",
+        "",
+        f"Question: {row['question']}",
+        "",
+        "Choices:",
+    ]
+    for choice in row["choices"]:
+        lines.append(f"  {choice['label']}. {choice['text']}")
+    lines += [
+        "",
+        "Answer (single letter only):",
+    ]
+    return "\n".join(lines)
+
+
+def _parse_mc_answer(response: str) -> str | None:
+    """Extract the first A/B/C/D letter from a model response."""
+    match = re.search(r"\b([A-D])\b", response.strip())
+    return match.group(1) if match else None
+
+
+def _benchmark_general_knowledge_multiple_choice(
+    runner,
+    model: str,
+    num_proc: int,
+    out_path: Path | None,
+) -> dict:
+    config = "general_knowledge_multiple_choice"
+    print(f"\n[{config}] Loading dataset...")
+    train_data = load_dataset(DATASET_NAME, config, num_proc=num_proc)["train"]
+
+    prompts = [_build_mc_prompt(row) for row in train_data]
+    responses = runner(prompts, MAX_NEW_TOKENS_MULTIPLE_CHOICE)
+
+    results = []
+    correct_count = 0
+    no_response_count = 0
+
+    for i, row in enumerate(train_data):
+        predicted = _parse_mc_answer(responses[i])
+        correct = row["correct_answer"]
+        is_correct = predicted == correct if predicted is not None else False
+        if predicted is None:
+            no_response_count += 1
+        if is_correct:
+            correct_count += 1
+
+        results.append(
+            {
+                "source": row["source"],
+                "process": row["process"],
+                "question": row["question"],
+                "choices": row["choices"],
+                "correct_answer": correct,
+                "response": responses[i],
+                "predicted": predicted,
+                "is_correct": is_correct,
+            }
+        )
+
+    total = len(results)
+    answered = total - no_response_count
+    accuracy = round(correct_count / total, 4) if total else 0.0
+
+    report = {
+        "model": model,
+        "dataset": DATASET_NAME,
+        "config": config,
+        "timestamp": datetime.now(timezone.utc).isoformat(),
+        "total_questions": total,
+        "no_response": no_response_count,
+        "answered": answered,
+        "correct": correct_count,
+        "accuracy": accuracy,
+        "results": results,
+    }
+
+    _print_gkmc_report(report)
+
+    if out_path is not None:
+        report_file = Path(out_path) / f"{config}.json"
+        with open(report_file, "w") as f:
+            json.dump(report, f, indent=2)
+        print(f"Report saved to: {report_file}")
+
+    return report
+
+
+def _print_gkmc_report(report: dict):
+    sep = "-" * 60
+    print(sep)
+    print("General Knowledge Multiple Choice — Report")
+    print(sep)
+    print(f"Model:            {report['model']}")
+    print(f"Total questions:  {report['total_questions']}")
+    print(f"No response:      {report['no_response']}")
+    print(f"Answered:         {report['answered']}")
+    print(f"Correct:          {report['correct']}")
+    print(f"Accuracy:         {report['accuracy']:.2%}")
+    print(sep)
+    for i, r in enumerate(report["results"], 1):
+        mark = "✓" if r["is_correct"] else "✗"
+        print(f"[{i:>3}] {mark} {r['source']} ({r['process']})")
+        print(f"       Q: {r['question'][:100]}")
+        print(f"       Predicted: {r['predicted']}  Correct: {r['correct_answer']}")
+    print(sep)
+    print(
+        f"Accuracy: {report['accuracy']:.2%}  ({report['correct']}/{report['total_questions']})"
+    )
+    print(sep)