Merge pull request #9 from ppak10/v0.0.21

V0.0.21

Author: ppak10

.gitignore

@@ -175,3 +175,6 @@ workspaces
 .claude/
 .gemini/
 .vscode/
+
+junit.xml
+junitxml

README.md

@@ -8,7 +8,7 @@
 
 Additive Manufacturing related software modules
 <p align="center">
-  <img src="./icon.svg" alt="Logo" width="50%">
+  <img src="./public/logo.svg" alt="Logo" width="50%">
 </p>
 
 ## Getting Started
@@ -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>
+```

icon.svg

@@ -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",
 ]

junit.xml

@@ -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

junitxml

@@ -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