diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e0322c0..f348314 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -14,6 +14,19 @@ permissions: contents: read jobs: + smoke-cpu: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: jdx/mise-action@v2 + - uses: actions/cache@v4 + with: + path: ~/.cache/uv + key: uv-${{ runner.os }}-${{ hashFiles('uv.lock') }} + - run: uv sync --extra training + - run: uv run python -m brainforge.training.smoke_cpu + env: + HF_HOME: ${{ runner.temp }}/hf-cache lint: name: Lint (ruff) runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index aa9e283..4453ea0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,63 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Task-level evaluation harness (`train task-eval`): verdict accuracy, FP/FN + rates and CWE classification accuracy on the post-cutoff holdout, scored + against the judge-verified answer stored in each record (`task_eval.json`). +- Interactive chat with a trained student: `train chat` (adapter or merged + model, greedy decoding via the model chat template). +- Training checkpointing: `save_strategy="steps"` with `TrainingConfig.save_steps` + (default 100, keep last 2), seeded runs (`TrainingConfig.seed`, default 42) + and `train run --resume` (TRL native resume from the latest checkpoint). +- CPU smoke training job in CI: 1 LoRA step on a tiny model, no GPU or + bitsandbytes needed (`make train-smoke-cpu`, `smoke-cpu` CI job). The smoke + test skips only when torch is missing; missing-training-deps or an + unreachable HF hub now fail the job instead of silently passing. +- `train evaluate`/`train export`/`train task-eval`/`train chat` now default + `--model` to the latest run directory under `experiments/`. + +### Changed + +- `evaluate` now honors the configured quantization (`--quantization + 4bit|8bit|none`) instead of always loading in 4-bit. +- `train prepare` writes the default output to `datasets/prepared/`, the + directory `train run`/`evaluate`/`task-eval` read by default; the `--output` + override is unchanged. +- `train run --resume` resumes the latest existing run directory instead of + creating a new one. +- `latest_run_dir` resolves a run by modification time and honors the + configured `training.output_dir`. + +### Fixed + +- `train task-eval` (CLI path): real generation now flattens the message list + into a single content string, matching the typed `generate(str | list[str])` + contract; `prediction_from_text` no longer swallows unexpected errors as + "miss". +- `train run`/`evaluate`/`task-eval`/`chat`/`export` fail with a clear message + instead of a raw traceback when an input file or model path is missing + (OSError clamped to a clean exit). +- `train run --resume` on a run without checkpoints fails fast with a clear + message instead of a raw HF Trainer `ValueError`. +- Removed dead code in `train chat` (misleading `name_or_path` argument, + redundant `generate_reply_fn`) and the pointless double dataset read in + `dataset split`. +- `case.id` is validated against `[A-Za-z0-9_-]{1,64}`, preventing path + traversal via crafted case files. +- Post-cutoff records no longer leak into the `train` split: they are now + exclusive to `test_postcutoff` (the anti-contamination benchmark). +- Response cache is now consulted by `structured()` calls; previously every + pipeline structured call bypassed the sqlite cache. + +### Security + +- `ProviderConfig.base_url` now validates scheme (http/https) and rejects + link-local and cloud-metadata hosts (SSRF guard); loopback stays allowed for + local gateways. +- Removed the committed `MLGW_API_KEY:deadbeef` placeholder default from the + sample config and docs (empty default like the other providers; the provider + fails at call time when the key is unset). + - Real QLoRA training pipeline: `train_qlora` (TRL `SFTTrainer`, bitsandbytes NF4/8bit quantization, gradient checkpointing, `TrainingConfig`-driven), `evaluate` (loss + perplexity on the post-cutoff split, `eval.json`) and diff --git a/Makefile b/Makefile index 4db6d18..708d878 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: install test lint format secrets audit build ci train-smoke train-run config models rag dataset validate train evaluate clean +.PHONY: install test lint format secrets audit build ci train-smoke train-smoke-cpu train-run config models rag dataset validate train evaluate clean install: uv sync @@ -33,6 +33,11 @@ train-smoke: uv sync --extra training uv run python -m brainforge.training.smoke +# CPU smoke: tiny model, no GPU/bitsandbytes needed; skips gracefully offline. +train-smoke-cpu: + uv sync --extra training + uv run python -m brainforge.training.smoke_cpu + train-run: uv run brainforge train run diff --git a/config/config.json b/config/config.json index 359349a..4ab54f1 100644 --- a/config/config.json +++ b/config/config.json @@ -14,7 +14,7 @@ "mlgw": { "type": "mlgw", "base_url": "${MLGW_BASE_URL:http://localhost:8080/v1}", - "api_key": "${MLGW_API_KEY:deadbeef}" + "api_key": "${MLGW_API_KEY:}" }, "mock": { "type": "mock" diff --git a/config/schema.json b/config/schema.json index e7c2d2c..e78ddb9 100644 --- a/config/schema.json +++ b/config/schema.json @@ -349,6 +349,18 @@ "title": "Quantization", "type": "string" }, + "save_steps": { + "default": 100, + "minimum": 1, + "title": "Save Steps", + "type": "integer" + }, + "seed": { + "default": 42, + "minimum": 0, + "title": "Seed", + "type": "integer" + }, "output_dir": { "default": "experiments", "title": "Output Dir", @@ -415,6 +427,8 @@ "batch_size": 1, "gradient_accumulation": 16, "quantization": "4bit", + "save_steps": 100, + "seed": 42, "output_dir": "experiments" } } diff --git a/docs/reference/providers.md b/docs/reference/providers.md index 8881389..15c3e56 100644 --- a/docs/reference/providers.md +++ b/docs/reference/providers.md @@ -49,11 +49,11 @@ styles. ## MLGW ```json -{"type": "mlgw", "base_url": "${MLGW_BASE_URL:http://localhost:8080/v1}", "api_key": "${MLGW_API_KEY:deadbeef}"} +{"type": "mlgw", "base_url": "${MLGW_BASE_URL:http://localhost:8080/v1}", "api_key": "${MLGW_API_KEY:}"} ``` -Local OpenAI-compatible gateway (llama.cpp, Ollama, vLLM backends). The -`deadbeef` development key is intentionally fake. +Local OpenAI-compatible gateway (llama.cpp, Ollama, vLLM backends). With an +empty key the provider fails at call time; set `MLGW_API_KEY` to use it. ## Local diff --git a/src/brainforge/case.py b/src/brainforge/case.py index ec79f56..ffebcd0 100644 --- a/src/brainforge/case.py +++ b/src/brainforge/case.py @@ -1,4 +1,8 @@ -from pydantic import BaseModel, Field +import re + +from pydantic import BaseModel, Field, field_validator + +_CASE_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]{1,64}$") class CaseSource(BaseModel): @@ -22,6 +26,13 @@ class Case(BaseModel): input: CaseInput metadata: dict = Field(default_factory=dict) + @field_validator("id") + @classmethod + def _validate_id(cls, value: str) -> str: + if _CASE_ID_PATTERN.match(value) is None: + raise ValueError("case id must match [A-Za-z0-9_-]{1,64}") + return value + @property def language(self) -> str | None: return self.metadata.get("language") diff --git a/src/brainforge/cli/dataset_cmd.py b/src/brainforge/cli/dataset_cmd.py index 7eef3f8..9f4b09d 100644 --- a/src/brainforge/cli/dataset_cmd.py +++ b/src/brainforge/cli/dataset_cmd.py @@ -111,8 +111,6 @@ def split( """Split a dataset into train/validation/test (+ post-cutoff holdout).""" from brainforge.training.prepare import prepare - records = _load_records(path) - del records target = output_dir or path.parent / (path.stem + "_split") try: stats = prepare(path, target, train_ratio, val_ratio, seed) diff --git a/src/brainforge/cli/train_cmd.py b/src/brainforge/cli/train_cmd.py index c1c4d3a..d19b767 100644 --- a/src/brainforge/cli/train_cmd.py +++ b/src/brainforge/cli/train_cmd.py @@ -1,4 +1,5 @@ from datetime import datetime +from functools import wraps from pathlib import Path import typer @@ -11,7 +12,40 @@ ) +def _clamp_errors(label: str): + """Turn expected failures (BrainforgeError and OS errors) into a clean exit.""" + + def decorator(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + except (BrainforgeError, OSError) as exc: + err_console.print(f"[red]{label} failed:[/red] {exc}") + raise typer.Exit(code=1) from exc + + return wrapper + + return decorator + + +def latest_run_dir(base: Path | None = None) -> Path: + """Resolve the most recent training run directory, by modification time.""" + if base is None: + from brainforge.config import load_config + + base = Path(load_config().training.output_dir) + base = Path(base).resolve() + candidates = [d for d in base.glob("run-*") if d.is_dir()] + if not candidates: + raise BrainforgeError( + f"no training runs found under {base}; run 'brainforge train run' first" + ) + return max(candidates, key=lambda d: d.stat().st_mtime) + + @app.command("prepare") +@_clamp_errors("prepare") def prepare( dataset: Path = typer.Argument(..., help="Path to dataset.jsonl"), output_dir: Path = typer.Option(None, "--output", "-o"), @@ -22,18 +56,15 @@ def prepare( """Validate, split and export a dataset for TRL training.""" from brainforge.training.prepare import prepare - target = output_dir or Path("datasets") / (dataset.stem + "_prepared") - try: - stats = prepare(dataset, target, train_ratio, val_ratio, seed) - except BrainforgeError as exc: - err_console.print(f"[red]prepare failed:[/red] {exc}") - raise typer.Exit(code=1) from exc + target = output_dir or Path("datasets") / "prepared" + stats = prepare(dataset, target, train_ratio, val_ratio, seed) console.print(f"[green]prepared dataset[/green] -> {target}") for name, count in stats.items(): console.print(f" {name}: {count}") @app.command("run") +@_clamp_errors("training") def run( dataset_dir: Path = typer.Option( Path("datasets/prepared"), @@ -51,18 +82,18 @@ def run( "-o", help="Output dir (default: /)", ), + resume: bool = typer.Option( + False, + "--resume/--no-resume", + help="Resume from the latest checkpoint in the output dir", + ), config: str = typer.Option(None, "--config", "-c"), ): """Run QLoRA training (requires a CUDA GPU and the training extra).""" from brainforge.config import load_config - from brainforge.errors import ConfigError from brainforge.training.qlora import train_qlora - try: - cfg = load_config(config) - except ConfigError as exc: - err_console.print(f"[red]config error:[/red] {exc}") - raise typer.Exit(code=1) from exc + cfg = load_config(config) training = cfg.training updates = {} if epochs: @@ -71,22 +102,22 @@ def run( updates["base_model"] = base_model if updates: training = training.model_copy(update=updates) - run_name = datetime.now().strftime("run-%Y%m%d-%H%M%S") - output_dir = output or Path(training.output_dir) / run_name - try: - summary = train_qlora(training, dataset_dir, output_dir) - except BrainforgeError as exc: - err_console.print(f"[red]training failed:[/red] {exc}") - raise typer.Exit(code=1) from exc + run_name = datetime.now().strftime("run-%Y%m%d-%H%M%S-%f") + if resume and output is None: + output_dir = latest_run_dir(base=Path(training.output_dir)) + else: + output_dir = output or Path(training.output_dir) / run_name + summary = train_qlora(training, dataset_dir, output_dir, resume=resume) console.print(f"[green]training done[/green] -> {summary['output_dir']}") for key, value in summary.items(): console.print(f" {key}: {value}") @app.command("evaluate") +@_clamp_errors("evaluation") def evaluate( model: Path = typer.Option( - Path("experiments"), "--model", "-m", help="Trained adapter directory" + None, "--model", "-m", help="Trained adapter directory (default: latest run)" ), dataset: Path = typer.Option( Path("datasets/prepared/test_postcutoff.jsonl"), @@ -98,20 +129,58 @@ def evaluate( """Evaluate a trained student (loss + perplexity, written to eval.json).""" from brainforge.training.qlora import evaluate as evaluate_model - try: - result = evaluate_model(model, dataset) - except BrainforgeError as exc: - err_console.print(f"[red]evaluation failed:[/red] {exc}") - raise typer.Exit(code=1) from exc + result = evaluate_model(model or latest_run_dir(), dataset) console.print(f"[green]evaluated[/green] {result['n_records']} records") console.print(f" eval_loss: {result['eval_loss']:.4f}") console.print(f" perplexity: {result['perplexity']:.4f}") +@app.command("chat") +@_clamp_errors("chat") +def chat( + model: Path = typer.Option( + None, "--model", "-m", help="Trained adapter or merged model (default: latest run)" + ), + max_new_tokens: int = typer.Option(512, "--max-new-tokens"), +): + """Interactive chat with a trained student model (requires CUDA + training extra).""" + from brainforge.training.chat import chat_loop + + chat_loop(model or latest_run_dir(), max_new_tokens=max_new_tokens) + + +@app.command("task-eval") +@_clamp_errors("task evaluation") +def task_eval( + model: Path = typer.Option( + None, "--model", "-m", help="Trained adapter or merged model (default: latest run)" + ), + dataset: Path = typer.Option( + Path("datasets/prepared/test_postcutoff.jsonl"), + "--dataset", + "-d", + help="Eval split (JSONL)", + ), + quantization: str = typer.Option("4bit", "--quantization", help="4bit, 8bit or none"), + max_new_tokens: int = typer.Option(512, "--max-new-tokens"), +): + """Task-level evaluation: verdict + CWE accuracy on a held-out split.""" + from brainforge.training.task_eval import evaluate_model_on_records + + result = evaluate_model_on_records( + model or latest_run_dir(), dataset, quantization, max_new_tokens + ) + console.print(f"[green]task evaluation done[/green] ({result['n_records']} records)") + for key in ("accuracy", "false_positive_rate", "false_negative_rate", "cwe_accuracy"): + value = result.get(key) + console.print(f" {key}: {value:.4f}" if isinstance(value, float) else f" {key}: {value}") + + @app.command("export") +@_clamp_errors("export") def export( model: Path = typer.Option( - Path("experiments"), "--model", "-m", help="Trained adapter directory" + None, "--model", "-m", help="Trained adapter directory (default: latest run)" ), output: Path = typer.Option( Path("models/export"), "--output", "-o", help="Merged model output directory" @@ -120,9 +189,5 @@ def export( """Merge the LoRA adapter into the base model and export it standalone.""" from brainforge.training.qlora import export as export_model - try: - result = export_model(model, output) - except BrainforgeError as exc: - err_console.print(f"[red]export failed:[/red] {exc}") - raise typer.Exit(code=1) from exc + result = export_model(model or latest_run_dir(), output) console.print(f"[green]exported[/green] -> {result['output_dir']}") diff --git a/src/brainforge/config/models.py b/src/brainforge/config/models.py index ead818d..868fb90 100644 --- a/src/brainforge/config/models.py +++ b/src/brainforge/config/models.py @@ -1,4 +1,5 @@ import re +import urllib.parse from typing import Literal from pydantic import BaseModel, ConfigDict, Field, field_validator @@ -6,6 +7,11 @@ from brainforge.types import ApiStyle, Mode, ProviderType, RoleKind _KNOWLEDGE_CUTOFF_RE = re.compile(r"^\d{4}-(0[1-9]|1[0-2])$") +_BLOCKED_BASE_URL_HOSTS = ( + "0.0.0.0", + "metadata.google.internal", +) +_BLOCKED_BASE_URL_PREFIXES = ("169.254.", "fe80:", "fc00:", "fd00:") class ProviderConfig(BaseModel): @@ -18,6 +24,21 @@ class ProviderConfig(BaseModel): timeout: float = Field(default=120.0, gt=0) max_retries: int = Field(default=3, ge=0) + @field_validator("base_url") + @classmethod + def _validate_base_url(cls, value: str | None) -> str | None: + if value is None: + return None + parsed = urllib.parse.urlparse(value) + if parsed.scheme not in ("http", "https"): + raise ValueError(f"base_url must use http or https, got '{parsed.scheme}'") + host = parsed.hostname + if not host: + raise ValueError(f"base_url must include a host: '{value}'") + if host in _BLOCKED_BASE_URL_HOSTS or host.startswith(_BLOCKED_BASE_URL_PREFIXES): + raise ValueError(f"base_url host '{host}' is not allowed (link-local or metadata)") + return value + class PricingConfig(BaseModel): model_config = ConfigDict(extra="forbid") @@ -77,6 +98,8 @@ class TrainingConfig(BaseModel): batch_size: int = Field(default=1, ge=1) gradient_accumulation: int = Field(default=16, ge=1) quantization: Literal["4bit", "8bit", "none"] = "4bit" + save_steps: int = Field(default=100, ge=1) + seed: int = Field(default=42, ge=0) output_dir: str = "experiments" diff --git a/src/brainforge/dataset/split.py b/src/brainforge/dataset/split.py index 20e1012..77aaa21 100644 --- a/src/brainforge/dataset/split.py +++ b/src/brainforge/dataset/split.py @@ -55,9 +55,9 @@ def split_with_postcutoff( train_ratio: float = 0.8, val_ratio: float = 0.1, seed: int = 42, - min_postcutoff: int = 20, ) -> dict[str, list[DatasetRecord]]: - splits = split_dataset(records, train_ratio, val_ratio, seed) + precutoff = [record for record in records if not is_postcutoff(record)] + splits = split_dataset(precutoff, train_ratio, val_ratio, seed) splits["test_postcutoff"] = postcutoff_records(records) return splits diff --git a/src/brainforge/providers/cache.py b/src/brainforge/providers/cache.py index 9c91133..da8aef4 100644 --- a/src/brainforge/providers/cache.py +++ b/src/brainforge/providers/cache.py @@ -4,8 +4,6 @@ import time from pathlib import Path -from pydantic import BaseModel - from brainforge.providers.base import ChatRequest, ChatResponse, Provider @@ -29,9 +27,6 @@ def complete(self, request: ChatRequest, model: str) -> ChatResponse: self._store(key, response) return response - def structured(self, request: ChatRequest, model: str, schema: type[BaseModel]) -> ChatResponse: - return self.inner.structured(request, model, schema) - def _key(self, request: ChatRequest, model: str) -> str: payload = json.dumps( { diff --git a/src/brainforge/training/chat.py b/src/brainforge/training/chat.py new file mode 100644 index 0000000..84e7320 --- /dev/null +++ b/src/brainforge/training/chat.py @@ -0,0 +1,80 @@ +"""Interactive chat with a trained (adapter or merged) student model.""" + +from pathlib import Path + +from brainforge.errors import BrainforgeError + +_EXIT_COMMANDS = {"quit", "exit"} + + +def load_chat_model(model_path): + """Load a merged model or a LoRA adapter directory for interactive chat.""" + model_path = Path(model_path) + if not model_path.exists(): + raise BrainforgeError(f"model directory not found: {model_path}") + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(str(model_path)) + if tokenizer.chat_template is None: + raise BrainforgeError(f"tokenizer at {model_path} has no chat template") + if (model_path / "adapter_config.json").is_file(): + from peft import AutoPeftModelForCausalLM + + model = AutoPeftModelForCausalLM.from_pretrained( + str(model_path), device_map="auto", torch_dtype=torch.bfloat16 + ) + else: + model = AutoModelForCausalLM.from_pretrained( + str(model_path), device_map="auto", torch_dtype=torch.bfloat16 + ) + model.eval() + return model, tokenizer + + +def generate_reply(model, tokenizer, history: list[str], max_new_tokens: int = 512) -> str: + """Generate one greedy assistant reply for the flat user-turn history.""" + import torch + + messages = [{"role": "user", "content": turn} for turn in history] + inputs = tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ).to(model.device) + with torch.no_grad(): + output = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) + generated = output[0][inputs["input_ids"].shape[1] :] + return tokenizer.decode(generated, skip_special_tokens=True).strip() + + +def chat_loop( + model_path, + reply_fn=None, + input_fn=input, + print_fn=print, + max_new_tokens: int = 512, +) -> None: + """Interactive REPL; exits on quit/exit/EOF, ignores blank lines.""" + if reply_fn is None: + model, tokenizer = load_chat_model(model_path) + reply_fn = lambda history: generate_reply( # noqa: E731 + model, tokenizer, history, max_new_tokens + ) + history: list[str] = [] + print_fn("brainforge chat - type 'quit' or 'exit' to leave") + while True: + try: + user_input = input_fn("you> ") + except EOFError: + print_fn("") + return + text = user_input.strip() + if not text: + continue + if text.lower() in _EXIT_COMMANDS: + return + history.append(text) + print_fn(f"assistant> {reply_fn(history)}") diff --git a/src/brainforge/training/qlora.py b/src/brainforge/training/qlora.py index 7caa7c1..1ad522c 100644 --- a/src/brainforge/training/qlora.py +++ b/src/brainforge/training/qlora.py @@ -48,18 +48,18 @@ def _messages_dataset(split_path: Path): return Dataset.from_list(extract_messages(read_jsonl(split_path))) -def _quantization_config(config: TrainingConfig): +def _quantization_config(quantization: str): import torch from transformers import BitsAndBytesConfig - if config.quantization == "4bit": + if quantization == "4bit": return BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, ) - if config.quantization == "8bit": + if quantization == "8bit": return BitsAndBytesConfig(load_in_8bit=True) return None @@ -76,14 +76,24 @@ def _lora_config(config: TrainingConfig): ) -def train_qlora(config: TrainingConfig, dataset_dir, output_dir) -> dict: +def train_qlora( + config: TrainingConfig, dataset_dir, output_dir, resume: bool | str = False +) -> dict: """Run QLoRA fine-tuning on a prepared dataset and save the adapter.""" + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + if resume: + checkpoints = sorted(output_dir.glob("checkpoint-*")) + if not checkpoints: + raise BrainforgeError( + f"no checkpoint found in {output_dir} to resume from;" + " run 'brainforge train run' without --resume first" + ) + resume = checkpoints[-1] _require_cuda() from trl import SFTConfig, SFTTrainer dataset_dir = Path(dataset_dir) - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) train_dataset = _messages_dataset(dataset_dir / "train.jsonl") validation_path = dataset_dir / "validation.jsonl" eval_dataset = _messages_dataset(validation_path) if validation_path.exists() else None @@ -99,7 +109,10 @@ def train_qlora(config: TrainingConfig, dataset_dir, output_dir) -> dict: logging_steps=1, eval_strategy="steps" if eval_dataset else "no", eval_steps=10, - save_strategy="no", + save_strategy="steps", + save_steps=config.save_steps, + save_total_limit=2, + seed=config.seed, report_to=[], ) trainer = SFTTrainer( @@ -107,10 +120,10 @@ def train_qlora(config: TrainingConfig, dataset_dir, output_dir) -> dict: args=sft_args, train_dataset=train_dataset, eval_dataset=eval_dataset, - quantization_config=_quantization_config(config), + quantization_config=_quantization_config(config.quantization), peft_config=_lora_config(config), ) - trainer.train() + trainer.train(resume_from_checkpoint=resume or None) trainer.save_model(str(output_dir)) summary = { "base_model": config.base_model, @@ -127,7 +140,7 @@ def train_qlora(config: TrainingConfig, dataset_dir, output_dir) -> dict: return summary -def evaluate(model_path, eval_dataset) -> dict: +def evaluate(model_path, eval_dataset, quantization: str = "4bit") -> dict: """Compute eval loss and perplexity of a trained adapter on a dataset split.""" _require_cuda() import torch @@ -143,12 +156,11 @@ def evaluate(model_path, eval_dataset) -> dict: raise BrainforgeError( f"tokenizer at {model_path} has no chat template; cannot evaluate chat records" ) - model = AutoPeftModelForCausalLM.from_pretrained( - str(model_path), - device_map="auto", - torch_dtype=torch.bfloat16, - load_in_4bit=True, - ) + load_kwargs = {"device_map": "auto", "torch_dtype": torch.bfloat16} + quant_config = _quantization_config(quantization) + if quant_config is not None: + load_kwargs["quantization_config"] = quant_config + model = AutoPeftModelForCausalLM.from_pretrained(str(model_path), **load_kwargs) model.eval() losses = [] with torch.no_grad(): diff --git a/src/brainforge/training/smoke_cpu.py b/src/brainforge/training/smoke_cpu.py new file mode 100644 index 0000000..115c555 --- /dev/null +++ b/src/brainforge/training/smoke_cpu.py @@ -0,0 +1,35 @@ +"""CPU-only smoke training: 1 LoRA step on a tiny model, CI-safe without GPU/HF.""" + +import sys + + +def main() -> int: + try: + import torch # noqa: F401 + except ImportError: + print("torch not installed; skipping CPU training smoke test") + return 0 + try: + from peft import LoraConfig, get_peft_model + from transformers import AutoModelForCausalLM, AutoTokenizer + + model_name = "hf-internal-testing/tiny-random-LlamaForCausalLM" + print(f"smoke test (CPU): 1 LoRA step on {model_name}") + tokenizer = AutoTokenizer.from_pretrained(model_name) + model = AutoModelForCausalLM.from_pretrained(model_name) + lora = LoraConfig(r=4, lora_alpha=8, lora_dropout=0.05, task_type="CAUSAL_LM") + model = get_peft_model(model, lora) + inputs = tokenizer("subprocess.run(user_input, shell=True)", return_tensors="pt").to( + model.device + ) + outputs = model(**inputs, labels=inputs["input_ids"]) + outputs.loss.backward() + print(f"smoke test passed (CPU, fp32): loss={outputs.loss.item():.4f}") + return 0 + except Exception as exc: + print(f"smoke test failed: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/brainforge/training/task_eval.py b/src/brainforge/training/task_eval.py new file mode 100644 index 0000000..abc7703 --- /dev/null +++ b/src/brainforge/training/task_eval.py @@ -0,0 +1,162 @@ +"""Task-level evaluation of a fine-tuned student model on a held-out split. + +Unlike ``qlora.evaluate`` (loss + perplexity), this harness scores the model's +actual task output (verdict + CWE classification) against the judge-verified +ground truth stored in each dataset record's assistant message. +""" + +import json +from pathlib import Path + +from brainforge.errors import BrainforgeError, ProviderError +from brainforge.providers.base import extract_json + + +def expected_from_record(record: dict) -> dict: + """Extract {vulnerability_found, cwe} from the record's judge answer.""" + messages = record.get("messages") or [] + if not messages: + raise BrainforgeError(f"record '{record.get('id')}' has no messages") + content = messages[-1].get("content", "") + try: + data = extract_json(content) + except Exception as exc: + raise BrainforgeError( + f"record '{record.get('id')}' has unparseable assistant JSON: {exc}" + ) from exc + return { + "vulnerability_found": data.get("verdict") == "confirmed", + "cwe": data.get("cwe"), + } + + +def prediction_from_text(text: str) -> dict: + """Map a raw model answer to {vulnerability_found, cwe}; None = unparseable miss.""" + try: + data = extract_json(text) + except ProviderError: + return {"vulnerability_found": None, "cwe": None} + return { + "vulnerability_found": data.get("verdict") == "confirmed", + "cwe": data.get("cwe"), + } + + +def score_task(pairs: list[tuple[dict, dict]]) -> dict: + """Compute accuracy, FP/FN rates and CWE accuracy from (prediction, expected) pairs.""" + if not pairs: + raise BrainforgeError("no evaluation pairs") + n = len(pairs) + correct = 0 + false_positives = 0 + false_negatives = 0 + cwe_correct = 0 + cwe_compared = 0 + for prediction, expected in pairs: + predicted = prediction["vulnerability_found"] + actual = expected["vulnerability_found"] + if predicted is not None and predicted == actual: + correct += 1 + elif (predicted is True and actual is False) or (predicted is None and actual is False): + false_positives += 1 + elif (predicted is False and actual is True) or (predicted is None and actual is True): + false_negatives += 1 + expected_cwe = expected.get("cwe") + if actual and expected_cwe: + cwe_compared += 1 + if prediction.get("cwe") == expected_cwe: + cwe_correct += 1 + return { + "n_records": n, + "accuracy": correct / n, + "false_positive_rate": false_positives / n, + "false_negative_rate": false_negatives / n, + "cwe_accuracy": (cwe_correct / cwe_compared) if cwe_compared else None, + } + + +def run_task_eval(records: list[dict], generate, max_new_tokens: int = 512) -> dict: + """Score generated answers against expected ones. + + ``generate`` receives the user message content and returns raw model text; + pass ``None`` only from the CLI path where the real model loader builds it. + """ + if not records: + raise BrainforgeError("empty evaluation dataset") + domains = {record.get("domain") for record in records} + unexpected = domains - {"security", None} + if unexpected: + raise BrainforgeError( + f"task evaluation only supports the security domain, found: {sorted(unexpected)}" + ) + pairs = [] + for record in records: + user_messages = [m["content"] for m in record.get("messages", []) if m["role"] == "user"] + text = generate(user_messages) + pairs.append((prediction_from_text(text), expected_from_record(record))) + result = score_task(pairs) + result["domains"] = sorted(d for d in domains if d) + return result + + +def evaluate_model_on_records( + model_path, dataset_path, quantization: str = "4bit", max_new_tokens: int = 512 +) -> dict: + """Load the trained model and run task evaluation on a JSONL dataset split.""" + from brainforge.dataset.writer import read_jsonl + from brainforge.training.qlora import _require_cuda + + _require_cuda() + model_path = Path(model_path) + dataset_path = Path(dataset_path) + records = read_jsonl(dataset_path) + generate = _build_generator(model_path, quantization, max_new_tokens) + result = run_task_eval(records, generate) + (model_path / "task_eval.json").write_text(json.dumps(result, indent=2), encoding="utf-8") + return result + + +def _flatten_message(user_message: str | list[str]) -> str: + """Normalize a user message (possibly a chat-style list) to a single string.""" + if isinstance(user_message, list): + return "\n\n".join(user_message) + return user_message + + +def _build_generator(model_path: Path, quantization: str, max_new_tokens: int): + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + + from brainforge.training.qlora import _quantization_config + + if not model_path.exists(): + raise BrainforgeError(f"model directory not found: {model_path}") + tokenizer = AutoTokenizer.from_pretrained(str(model_path)) + if tokenizer.chat_template is None: + raise BrainforgeError(f"tokenizer at {model_path} has no chat template") + load_kwargs = {"device_map": "auto", "torch_dtype": torch.bfloat16} + quant_config = _quantization_config(quantization) + if quant_config is not None: + load_kwargs["quantization_config"] = quant_config + if (model_path / "adapter_config.json").is_file(): + from peft import AutoPeftModelForCausalLM + + model = AutoPeftModelForCausalLM.from_pretrained(str(model_path), **load_kwargs) + else: + model = AutoModelForCausalLM.from_pretrained(str(model_path), **load_kwargs) + model.eval() + + def generate(user_message: str | list[str]) -> str: + inputs = tokenizer.apply_chat_template( + [{"role": "user", "content": _flatten_message(user_message)}], + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ).to(model.device) + with torch.no_grad(): + output = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) + generated = output[0][inputs["input_ids"].shape[1] :] + return tokenizer.decode(generated, skip_special_tokens=True) + + return generate diff --git a/tests/unit/test_chat.py b/tests/unit/test_chat.py new file mode 100644 index 0000000..3f56e9d --- /dev/null +++ b/tests/unit/test_chat.py @@ -0,0 +1,80 @@ +import pytest + +from brainforge.training.chat import chat_loop, generate_reply, load_chat_model + + +class TestChatLoop: + def test_roundtrip(self): + history = [] + inputs = iter(["hello", "quit"]) + printed = [] + chat_loop( + "unused-model-path", + reply_fn=lambda h: history.append(h) or "ok", + input_fn=lambda prompt: next(inputs), + print_fn=lambda text: printed.append(text), + ) + # reply_fn receives the conversation, loop keeps full history + assert history == [["hello"]] + assert any("ok" in line for line in printed) + + def test_history_accumulates(self): + calls = [] + + def reply(history): + calls.append(list(history)) + return f"reply-{len(calls)}" + + inputs = iter(["one", "two", "quit"]) + chat_loop( + "unused", + reply_fn=reply, + input_fn=lambda prompt: next(inputs), + print_fn=lambda text: None, + ) + assert calls[0] == ["one"] + assert calls[1] == ["one", "two"] + + def test_exit_command(self): + inputs = iter(["exit"]) + chat_loop( + "unused", + reply_fn=lambda h: "x", + input_fn=lambda p: next(inputs), + print_fn=lambda t: None, + ) + + def test_eof_exits(self): + # input_fn raises EOFError like input() on Ctrl+D + def eof(prompt): + raise EOFError + + chat_loop("unused", reply_fn=lambda h: "x", input_fn=eof, print_fn=lambda t: None) + + def test_blank_lines_ignored(self): + replies = [] + + def reply(history): + replies.append(list(history)) + return "ok" + + inputs = iter(["", " ", "hello", "quit"]) + chat_loop( + "unused", reply_fn=reply, input_fn=lambda p: next(inputs), print_fn=lambda t: None + ) + assert replies == [["hello"]] + + +class TestLoadChatModel: + def test_missing_path_fails_cleanly(self, tmp_path): + from brainforge.errors import BrainforgeError + + with pytest.raises(BrainforgeError, match="not found"): + load_chat_model(tmp_path / "ghost") + + +def test_generate_reply_signature(): + import inspect + + params = inspect.signature(generate_reply).parameters + assert list(params) == ["model", "tokenizer", "history", "max_new_tokens"] diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index ad206b8..66fb6aa 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -109,10 +109,8 @@ def test_train_prepare(project_env): ) result = runner.invoke(app, ["train", "prepare", "datasets/security_dataset.jsonl"]) assert result.exit_code == 0, result.output - assert (project_env / "datasets" / "security_dataset_prepared" / "train.jsonl").is_file() - assert ( - project_env / "datasets" / "security_dataset_prepared" / "test_postcutoff.jsonl" - ).is_file() + assert (project_env / "datasets" / "prepared" / "train.jsonl").is_file() + assert (project_env / "datasets" / "prepared" / "test_postcutoff.jsonl").is_file() def test_train_run_fails_without_gpu(project_env): @@ -121,6 +119,18 @@ def test_train_run_fails_without_gpu(project_env): assert "training failed" in result.output +def test_train_evaluate_no_runs_fails_cleanly(project_env): + result = runner.invoke(app, ["train", "evaluate"]) + assert result.exit_code == 1 + assert "no training runs" in result.output + + +def test_train_export_no_runs_fails_cleanly(project_env): + result = runner.invoke(app, ["train", "export"]) + assert result.exit_code == 1 + assert "no training runs" in result.output + + def test_rag_index_and_search_hashing(project_env): result = runner.invoke(app, ["rag", "index", "data/raw", "--backend", "hashing"]) assert result.exit_code == 0, result.output diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 588f42c..8ba8cce 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -204,6 +204,34 @@ def test_knowledge_cutoff_format(config_file): Config.model_validate(make_raw() | {"models": raw["models"]}) +def test_provider_base_url_rejects_metadata_and_link_local(): + from brainforge.config.models import ProviderConfig + + for bad in ( + "http://169.254.169.254/latest/meta-data/", + "http://metadata.google.internal/computeMetadata/", + "https://0.0.0.0/v1", + "ftp://openrouter.ai/v1", + ): + with pytest.raises(ValidationError, match="base_url"): + ProviderConfig(type="mock", base_url=bad) + for good in ( + "https://openrouter.ai/api/v1", + "http://localhost:8080/v1", + "http://127.0.0.1:11434/v1", + ): + ProviderConfig(type="mock", base_url=good) + + +def test_default_config_mlgw_key_has_no_fake_default(monkeypatch): + from brainforge.config.defaults import DEFAULT_CONFIG_PATH + + monkeypatch.delenv("MLGW_API_KEY", raising=False) + expanded = expand_env(json.loads(DEFAULT_CONFIG_PATH.read_text())) + api_key = expanded["providers"]["mlgw"]["api_key"] + assert api_key != "deadbeef" + + def test_schema_roundtrip(config_file, tmp_path): schema_path = tmp_path / "schema.json" write_schema(schema_path) @@ -215,6 +243,16 @@ def test_schema_roundtrip(config_file, tmp_path): assert not schema_matches(schema_path) +def test_training_checkpoint_fields(): + from brainforge.config.models import TrainingConfig + + config = TrainingConfig() + assert config.save_steps == 100 + assert config.seed == 42 + with pytest.raises(ValidationError): + TrainingConfig(save_steps=0) + + def test_training_defaults(config_file): config = load_config(config_file) assert config.training.base_model == "Qwen/Qwen3-8B" diff --git a/tests/unit/test_pipeline_dataset.py b/tests/unit/test_pipeline_dataset.py index 5fc050c..9213c49 100644 --- a/tests/unit/test_pipeline_dataset.py +++ b/tests/unit/test_pipeline_dataset.py @@ -8,6 +8,7 @@ from brainforge.dataset.dedup import deduplicate from brainforge.dataset.split import ( group_key, + is_postcutoff, postcutoff_warning, split_dataset, split_with_postcutoff, @@ -305,6 +306,25 @@ def test_case_builder_from_dict(): assert case.source.date == "2026-08-01" +def test_case_id_rejects_path_traversal(): + from pydantic import ValidationError + + with pytest.raises(ValidationError): + Case( + id="../../etc/passwd", + source=CaseSource(type="manual"), + input=CaseInput(code="x"), + ) + with pytest.raises(ValidationError): + build_case_from_dict( + { + "id": "../../../tmp/pwn", + "source": {"type": "manual"}, + "input": {"code": "x"}, + } + ) + + def test_case_id_deterministic(): payload = {"source": {"type": "manual"}, "input": {"code": "x"}} assert build_case_from_dict(payload).id == build_case_from_dict(payload).id @@ -371,6 +391,33 @@ def record_for(case: Case) -> DatasetRecord: assert warning is not None and "post-cutoff" in warning +def test_postcutoff_records_excluded_from_regular_splits(): + def make_record(rid: str, repo: str, postcutoff: bool) -> DatasetRecord: + return DatasetRecord( + id=rid, + domain="security", + messages=[ + ChatMessage(role="user", content=rid), + ChatMessage(role="assistant", content="a"), + ], + metadata={ + "recitation_risk": not postcutoff, + "source": {"type": "git", "repository": repo}, + }, + ) + + # One post-cutoff record per repo so every hash bucket is covered. + records = [make_record(f"post-{i}", f"repo-post-{i}", True) for i in range(30)] + records += [make_record(f"pre-{i}", f"repo-pre-{i}", False) for i in range(30)] + splits = split_with_postcutoff(records) + for name in ("train", "validation", "test"): + leaked = [r.id for r in splits[name] if is_postcutoff(r)] + assert not leaked, f"{name} contains post-cutoff records: {leaked}" + assert len(splits["test_postcutoff"]) == 30 + total = sum(len(v) for v in splits.values()) + assert total == 60, "post-cutoff records must not be duplicated across splits" + + def test_dedup_exact_and_near(tmp_path): def make_record(rid: str, user: str, assistant: str) -> DatasetRecord: return DatasetRecord( diff --git a/tests/unit/test_providers.py b/tests/unit/test_providers.py index f17a18b..cbb0363 100644 --- a/tests/unit/test_providers.py +++ b/tests/unit/test_providers.py @@ -6,7 +6,7 @@ from brainforge.config.models import ProviderConfig from brainforge.errors import ProviderError, ProviderNotSupportedError, ProviderUnavailableError -from brainforge.providers.base import ChatMessage, ChatRequest, Provider, extract_json +from brainforge.providers.base import ChatMessage, ChatRequest, ChatResponse, Provider, extract_json from brainforge.providers.cache import CacheProvider from brainforge.providers.mock import MockProvider from brainforge.providers.observability import UsageLogger, estimate_cost @@ -82,7 +82,12 @@ def __init__(self): def complete(self, request, model): self.calls += 1 - return MockProvider(self.name, self.config).complete(request, model) + response = ChatResponse( + content=json.dumps({"verdict": "ok", "confidence": 0.5, "items": ["x"]}), + provider=self.name, + model=model, + ) + return response def test_extract_json_plain(): @@ -184,6 +189,18 @@ def test_cache_hit_and_miss(tmp_path): assert first.content == second.content +def test_cache_structured_hit_and_miss(tmp_path): + inner = CountingProvider() + cached = CacheProvider(inner, tmp_path / "cache.sqlite") + first = cached.structured(make_request(), "m", SimpleSchema) + second = cached.structured(make_request(), "m", SimpleSchema) + assert first.cached is False + assert second.cached is True + assert inner.calls == 1 + assert first.data is not None + assert second.data is not None + + def test_cache_disabled(tmp_path): inner = CountingProvider() cached = CacheProvider(inner, tmp_path / "cache.sqlite", enabled=False) diff --git a/tests/unit/test_task_eval.py b/tests/unit/test_task_eval.py new file mode 100644 index 0000000..118f5bb --- /dev/null +++ b/tests/unit/test_task_eval.py @@ -0,0 +1,113 @@ +import json + +import pytest + +from brainforge.errors import BrainforgeError +from brainforge.training.task_eval import ( + expected_from_record, + prediction_from_text, + run_task_eval, + score_task, +) + + +def make_record(verdict: str, cwe: str | None = "CWE-78", domain: str = "security") -> dict: + expected = {"verdict": verdict, "confidence": 0.9, "reasoning": "r", "evidence": []} + if cwe: + expected["cwe"] = cwe + return { + "id": f"rec-{verdict}", + "domain": domain, + "messages": [ + {"role": "user", "content": "analyze this snippet"}, + {"role": "assistant", "content": json.dumps(expected)}, + ], + "metadata": {}, + } + + +def test_expected_from_record_parses_judge_verdict(): + expected = expected_from_record(make_record("confirmed")) + assert expected == {"vulnerability_found": True, "cwe": "CWE-78"} + + +def test_expected_from_record_rejected(): + assert expected_from_record(make_record("rejected", cwe=None)) == { + "vulnerability_found": False, + "cwe": None, + } + + +def test_prediction_from_text_parses_json(): + text = '```json\n{"verdict": "confirmed", "cwe": "CWE-79"}\n```' + assert prediction_from_text(text) == {"vulnerability_found": True, "cwe": "CWE-79"} + + +def test_prediction_from_text_unparseable_is_miss(): + assert prediction_from_text("garbage, no json here") == { + "vulnerability_found": None, + "cwe": None, + } + + +def test_prediction_from_text_propagates_unexpected_errors(monkeypatch): + from brainforge.training import task_eval + + def boom(_text: str) -> dict: + raise ValueError("extract_json bug") + + monkeypatch.setattr(task_eval, "extract_json", boom) + with pytest.raises(ValueError): + prediction_from_text("x") + + +def test_generate_content_flattens_message_list(): + from brainforge.training.task_eval import _flatten_message + + assert _flatten_message("single") == "single" + assert _flatten_message(["first", "second"]) == "first\n\nsecond" + assert _flatten_message([]) == "" + + +def test_score_task_metrics(): + # 4 records: TP, TN, 1 FP, 1 FN + pairs = [ + ( + {"vulnerability_found": True, "cwe": "CWE-78"}, + {"vulnerability_found": True, "cwe": "CWE-78"}, + ), + ({"vulnerability_found": False, "cwe": None}, {"vulnerability_found": False, "cwe": None}), + ( + {"vulnerability_found": True, "cwe": "CWE-79"}, + {"vulnerability_found": False, "cwe": None}, + ), + ( + {"vulnerability_found": None, "cwe": None}, + {"vulnerability_found": True, "cwe": "CWE-78"}, + ), + ] + metrics = score_task(pairs) + assert metrics["n_records"] == 4 + assert metrics["accuracy"] == 0.5 + assert metrics["false_positive_rate"] == 0.25 + assert metrics["false_negative_rate"] == 0.25 + assert metrics["cwe_accuracy"] == 0.5 + + +def test_run_task_eval_rejects_non_security(): + records = [make_record("confirmed", domain="coding")] + with pytest.raises(BrainforgeError, match="domain"): + run_task_eval(records, generate=lambda messages: "{}") + + +def test_run_task_eval_with_stub_generate(): + records = [make_record("confirmed", cwe="CWE-78"), make_record("rejected", cwe=None)] + stub_reply = json.dumps({"verdict": "confirmed", "cwe": "CWE-78"}) + received = [] + result = run_task_eval( + records, generate=lambda messages: received.append(messages) or stub_reply + ) + assert result["accuracy"] == 0.5 + assert result["n_records"] == 2 + # One generate() call per record, user message only. + assert received == [["analyze this snippet"]] * 2 diff --git a/tests/unit/test_training_glue.py b/tests/unit/test_training_glue.py index 76a3135..9333768 100644 --- a/tests/unit/test_training_glue.py +++ b/tests/unit/test_training_glue.py @@ -1,4 +1,5 @@ import json +from pathlib import Path import pytest from typer.testing import CliRunner @@ -69,3 +70,121 @@ def test_cli_train_run_fails_cleanly_without_gpu(project_env, monkeypatch): result = runner.invoke(app, ["train", "run", "--dataset-dir", str(dataset_dir)]) assert result.exit_code == 1 assert "training failed" in result.output + + +def test_cli_train_run_resume_flag_fails_cleanly_without_gpu(project_env): + dataset_dir = project_env / "datasets" / "prepared" + dataset_dir.mkdir(parents=True) + split = [{"id": "x", "messages": [{"role": "user", "content": "hi"}], "metadata": {}}] + for name in ("train.jsonl", "validation.jsonl"): + (dataset_dir / name).write_text("\n".join(json.dumps(r) for r in split) + "\n") + result = runner.invoke(app, ["train", "run", "--dataset-dir", str(dataset_dir), "--resume"]) + assert result.exit_code == 1 + assert "training failed" in result.output + + +def test_latest_run_dir_uses_mtime_not_name(project_env): + import os + import time + + from brainforge.cli.train_cmd import latest_run_dir + + base = project_env / "experiments" + name_latest = base / "run-20260101-000000" + mtime_latest = base / "run-20251231-235959" + name_latest.mkdir(parents=True) + mtime_latest.mkdir(parents=True) + os.utime(name_latest, (0, 0)) + os.utime(mtime_latest, (time.time(), time.time())) + assert latest_run_dir().name == "run-20251231-235959" + + +def test_latest_run_dir_respects_config_output_dir(project_env): + from brainforge.cli.train_cmd import latest_run_dir + + (project_env / "experiments").mkdir(exist_ok=True) + run_dir = project_env / "artifacts" / "run-20260920-000001" + run_dir.mkdir(parents=True) + with (project_env / "config" / "config.json").open("r") as handle: + config = json.load(handle) + config["training"]["output_dir"] = "artifacts" + (project_env / "config" / "config.json").write_text(json.dumps(config)) + assert latest_run_dir().name == "run-20260920-000001" + + +def test_train_qlora_resume_without_checkpoint_fails_fast(tmp_path): + output_dir = tmp_path / "run" + with pytest.raises(BrainforgeError, match="no checkpoint"): + train_qlora(_training_config(), tmp_path, output_dir, resume=True) + + +def test_cli_train_run_resume_reuses_latest_run(project_env, monkeypatch): + calls = {} + + def fake_train_qlora(config, dataset_dir, output_dir, resume=False): + calls["output_dir"] = Path(output_dir) + calls["resume"] = resume + raise BrainforgeError("stop") + + monkeypatch.setattr("brainforge.training.qlora.train_qlora", fake_train_qlora) + latest = project_env / "experiments" / "run-20260920-000001" + latest.mkdir(parents=True) + result = runner.invoke(app, ["train", "run", "--resume"]) + assert result.exit_code == 1 + assert calls.get("output_dir") == latest + assert calls["resume"] is True + + +def test_cli_train_run_missing_dataset_fails_cleanly(project_env, monkeypatch): + import os + + def boom(config, dataset_dir, output_dir, resume=False): + raise FileNotFoundError(os.path.join(str(dataset_dir), "train.jsonl")) + + monkeypatch.setattr("brainforge.training.qlora.train_qlora", boom) + result = runner.invoke(app, ["train", "run"]) + assert result.exit_code == 1 + assert "training failed" in result.output + assert "Traceback" not in result.output + + +def test_cli_train_task_eval_fails_cleanly(project_env): + result = runner.invoke(app, ["train", "task-eval", "--model", str(project_env / "nope")]) + assert result.exit_code == 1 + assert "task evaluation failed" in result.output + + +def test_cli_train_chat_missing_model_fails(project_env): + result = runner.invoke(app, ["train", "chat", "--model", str(project_env / "nope")]) + assert result.exit_code == 1 + assert "chat failed" in result.output + + +def test_evaluate_signature_accepts_quantization(): + import inspect + + from brainforge.training.qlora import evaluate + + params = inspect.signature(evaluate).parameters + assert "quantization" in params + assert params["quantization"].default == "4bit" + + +def test_smoke_cpu_skips_without_torch(monkeypatch, capsys): + import sys + + monkeypatch.setitem(sys.modules, "torch", None) + from brainforge.training import smoke_cpu + + assert smoke_cpu.main() == 0 + assert "skipping" in capsys.readouterr().out.lower() + + +def test_smoke_cpu_fails_when_deps_missing(monkeypatch): + import sys + import types + + monkeypatch.setitem(sys.modules, "torch", types.ModuleType("torch")) + from brainforge.training import smoke_cpu + + assert smoke_cpu.main() == 1