mirror of https://github.com/razor-ai/soup.git
Add Phase 3: serve, data generate, sweep, diff, DeepSpeed (v0.3.0)
- soup serve: FastAPI inference server with OpenAI-compatible API, SSE streaming - soup data generate: synthetic data generation via OpenAI API or local models - soup sweep: grid/random hyperparameter search with experiment tracker integration - soup diff: side-by-side model comparison with metrics - Multi-GPU/DeepSpeed: ZeRO Stage 2/3 configs, --deepspeed flag in train command - 95 new tests (281 total), all passing - Removed TESTING_GUIDE.md Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Alpamys
parent
8b1844d316
commit
87fd760847
17
CLAUDE.md
17
CLAUDE.md
|
|
@ -61,6 +61,16 @@ soup train --config soup.yaml
|
|||
|
||||
**W&B integration:** `commands/train.py` supports `--wandb` flag to enable Weights & Biases logging. Sets `report_to="wandb"` in TrainingArguments. Requires `pip install wandb`.
|
||||
|
||||
**Serve:** `commands/serve.py` starts a local inference server with OpenAI-compatible API (`/v1/chat/completions`, `/v1/models`, `/health`). Uses FastAPI + uvicorn. Supports LoRA adapters and full models, SSE streaming. Requires `pip install 'soup-cli[serve]'`.
|
||||
|
||||
**Data generate:** `commands/generate.py` generates synthetic training data using LLMs. Supports OpenAI API and local models as providers. Outputs in alpaca/sharegpt/chatml format. Validates on the fly and can deduplicate against existing datasets. Registered as `soup data generate`.
|
||||
|
||||
**Sweep:** `commands/sweep.py` runs hyperparameter search (grid or random) over training parameters. Uses shortcut names (lr, epochs, lora_r, etc.) or dot notation. Each run is saved to the experiment tracker. Supports `--dry-run` to preview combinations.
|
||||
|
||||
**Diff:** `commands/diff.py` compares outputs of two models side-by-side on the same prompts. Computes metrics (length, word count, word overlap). Supports JSONL prompt files and CLI prompt arguments.
|
||||
|
||||
**DeepSpeed:** `utils/deepspeed.py` provides ZeRO Stage 2/3 config templates. `commands/train.py` supports `--deepspeed zero2|zero3|zero2_offload|<path>`. Trainers (SFT/DPO) pass `deepspeed` to HF TrainingArguments. Requires `pip install 'soup-cli[deepspeed]'`.
|
||||
|
||||
## Code Conventions
|
||||
|
||||
- **Line length:** 100 chars (ruff enforced)
|
||||
|
|
@ -89,7 +99,7 @@ soup train --config soup.yaml
|
|||
|
||||
## Tests
|
||||
|
||||
Test suite (~184 tests) lives in `tests/`:
|
||||
Test suite (~281 tests) lives in `tests/`:
|
||||
|
||||
| File | Covers |
|
||||
|---|---|
|
||||
|
|
@ -113,3 +123,8 @@ Test suite (~184 tests) lives in `tests/`:
|
|||
| `test_merge.py` | Merge command, adapter detection, validation |
|
||||
| `test_export.py` | Export command, GGUF quant types, validation |
|
||||
| `test_resume.py` | Resume checkpoint resolution, W&B flag |
|
||||
| `test_serve.py` | Serve command, FastAPI app, endpoints, streaming |
|
||||
| `test_generate.py` | Data generate, JSON parsing, validation, prompts |
|
||||
| `test_sweep.py` | Sweep params parsing, combinations, nested config |
|
||||
| `test_diff.py` | Diff prompts collection, metrics, CLI |
|
||||
| `test_deepspeed.py` | DeepSpeed configs, multi-GPU detection, trainer integration |
|
||||
|
|
|
|||
110
README.md
110
README.md
|
|
@ -21,7 +21,7 @@
|
|||
<a href="https://pypi.org/project/soup-cli/"><img src="https://img.shields.io/pypi/v/soup-cli?color=blue" alt="PyPI"></a>
|
||||
<img src="https://img.shields.io/badge/python-3.9%2B-blue" alt="Python 3.9+">
|
||||
<img src="https://img.shields.io/badge/license-MIT-green" alt="MIT License">
|
||||
<img src="https://img.shields.io/badge/tests-186%20passed-brightgreen" alt="Tests">
|
||||
<img src="https://img.shields.io/badge/tests-281%20passed-brightgreen" alt="Tests">
|
||||
<a href="https://github.com/MakazhanAlpamys/Soup/actions"><img src="https://github.com/MakazhanAlpamys/Soup/actions/workflows/ci.yml/badge.svg" alt="CI"></a>
|
||||
</p>
|
||||
|
||||
|
|
@ -228,6 +228,102 @@ soup train --config soup.yaml --wandb
|
|||
|
||||
Make sure `WANDB_API_KEY` is set or run `wandb login` first.
|
||||
|
||||
## Inference Server
|
||||
|
||||
Start a local OpenAI-compatible inference server:
|
||||
|
||||
```bash
|
||||
# Install server dependencies
|
||||
pip install 'soup-cli[serve]'
|
||||
|
||||
# Start server
|
||||
soup serve --model ./output --port 8000
|
||||
|
||||
# With custom settings
|
||||
soup serve --model ./output --port 8080 --host 127.0.0.1 --max-tokens 1024
|
||||
```
|
||||
|
||||
Endpoints:
|
||||
- `POST /v1/chat/completions` — chat completions (streaming supported)
|
||||
- `GET /v1/models` — list available models
|
||||
- `GET /health` — health check
|
||||
|
||||
Compatible with OpenAI SDK:
|
||||
```python
|
||||
from openai import OpenAI
|
||||
client = OpenAI(base_url="http://localhost:8000/v1", api_key="unused")
|
||||
response = client.chat.completions.create(
|
||||
model="output",
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
)
|
||||
```
|
||||
|
||||
## Synthetic Data Generation
|
||||
|
||||
Generate training data using LLMs:
|
||||
|
||||
```bash
|
||||
# Generate using OpenAI API
|
||||
soup data generate --prompt "Create math word problems" --count 100 --format alpaca
|
||||
|
||||
# Use a different model
|
||||
soup data generate --prompt "Medical Q&A pairs" --model gpt-4o --count 500
|
||||
|
||||
# Deduplicate against existing data
|
||||
soup data generate --prompt "..." --count 200 --dedup-with existing.jsonl
|
||||
|
||||
# Use seed examples to guide style
|
||||
soup data generate --prompt "..." --seed examples.jsonl --count 100
|
||||
```
|
||||
|
||||
## Hyperparameter Sweep
|
||||
|
||||
Search for the best hyperparameters:
|
||||
|
||||
```bash
|
||||
# Grid search over learning rate and LoRA rank
|
||||
soup sweep --config soup.yaml --param lr=1e-5,2e-5,5e-5 --param lora_r=8,16,32
|
||||
|
||||
# Random search with max runs
|
||||
soup sweep --config soup.yaml --param lr=1e-5,2e-5,5e-5 --strategy random --max-runs 5
|
||||
|
||||
# Preview without running
|
||||
soup sweep --config soup.yaml --param lr=1e-5,2e-5 --param epochs=2,3 --dry-run
|
||||
```
|
||||
|
||||
## Model Comparison
|
||||
|
||||
Compare outputs of two models side-by-side:
|
||||
|
||||
```bash
|
||||
# Compare with inline prompts
|
||||
soup diff --model-a ./model_v1 --model-b ./model_v2 --prompt "Explain gravity"
|
||||
|
||||
# Compare with a prompts file
|
||||
soup diff --model-a ./base --model-b ./finetuned --prompts test_prompts.jsonl
|
||||
|
||||
# Save results
|
||||
soup diff --model-a ./a --model-b ./b --prompts prompts.txt --output results.jsonl
|
||||
```
|
||||
|
||||
## Multi-GPU / DeepSpeed
|
||||
|
||||
Train on multiple GPUs with DeepSpeed:
|
||||
|
||||
```bash
|
||||
# ZeRO Stage 2 (recommended for most cases)
|
||||
soup train --config soup.yaml --deepspeed zero2
|
||||
|
||||
# ZeRO Stage 3 (for very large models)
|
||||
soup train --config soup.yaml --deepspeed zero3
|
||||
|
||||
# ZeRO Stage 2 with CPU offload (memory-constrained)
|
||||
soup train --config soup.yaml --deepspeed zero2_offload
|
||||
|
||||
# Custom DeepSpeed config
|
||||
soup train --config soup.yaml --deepspeed ./my_ds_config.json
|
||||
```
|
||||
|
||||
## Data Formats
|
||||
|
||||
Soup supports these formats (auto-detected):
|
||||
|
|
@ -328,6 +424,11 @@ soup eval --model ./output --benchmarks mmlu --run-id run_20260223_143052_a1b2
|
|||
| Experiment tracking (SQLite) | ✅ |
|
||||
| Data tools (convert, merge, dedup, stats) | ✅ |
|
||||
| Model evaluation (lm-eval) | ✅ |
|
||||
| Inference server (OpenAI-compatible) | ✅ |
|
||||
| Synthetic data generation | ✅ |
|
||||
| Hyperparameter sweep (grid/random) | ✅ |
|
||||
| Model comparison (diff) | ✅ |
|
||||
| Multi-GPU / DeepSpeed | ✅ |
|
||||
| Web dashboard | 🔜 |
|
||||
| Cloud mode (BYOG) | 🔜 |
|
||||
|
||||
|
|
@ -353,6 +454,11 @@ soup runs show <run_id> Detailed run info + loss graph
|
|||
soup runs compare <run_1> <run_2> Compare two runs
|
||||
soup runs delete <run_id> Remove a run
|
||||
soup eval --model ./output --benchmarks mmlu Evaluate on benchmarks
|
||||
soup serve --model ./output --port 8000 OpenAI-compatible API server
|
||||
soup sweep --config soup.yaml --param lr=... Hyperparameter search
|
||||
soup diff --model-a ./a --model-b ./b Compare two models
|
||||
soup data generate --prompt "..." --count 100 Generate synthetic data
|
||||
soup train --deepspeed zero2 Multi-GPU with DeepSpeed
|
||||
soup version Show version
|
||||
```
|
||||
|
||||
|
|
@ -372,7 +478,7 @@ pip install -e ".[dev]"
|
|||
# Lint
|
||||
ruff check soup_cli/ tests/
|
||||
|
||||
# Run unit tests (fast, no GPU needed — 186 tests)
|
||||
# Run unit tests (fast, no GPU needed — 281 tests)
|
||||
pytest tests/ -v
|
||||
|
||||
# Run smoke tests (downloads tiny model, runs real training)
|
||||
|
|
|
|||
196
TESTING_GUIDE.md
196
TESTING_GUIDE.md
|
|
@ -1,196 +0,0 @@
|
|||
# Soup CLI — Quick Local Test Guide (Windows)
|
||||
|
||||
**Hardware:** RTX 3050 (4GB VRAM), i5
|
||||
**Model:** `TinyLlama/TinyLlama-1.1B-Chat-v1.0` (~1.1B params, ~600MB in 4-bit)
|
||||
**OS:** Windows (CMD/PowerShell)
|
||||
|
||||
---
|
||||
|
||||
## 0. Install
|
||||
|
||||
```cmd
|
||||
pip install soup-cli
|
||||
pip install datasketch
|
||||
```
|
||||
|
||||
> Unit tests (`pytest tests/`) only work from the repo clone (`pip install -e ".[dev]"`), not from pip install.
|
||||
|
||||
## 1. Version & Help
|
||||
|
||||
```cmd
|
||||
soup version
|
||||
soup --help
|
||||
soup train --help
|
||||
soup data --help
|
||||
```
|
||||
|
||||
## 2. Init Config from Templates
|
||||
|
||||
```cmd
|
||||
soup init -t chat -o test_chat.yaml
|
||||
soup init -t code -o test_code.yaml
|
||||
```
|
||||
|
||||
> `soup init` without `-t` opens interactive wizard (requires terminal input).
|
||||
|
||||
## 3. Create Test Dataset
|
||||
|
||||
Create file `test_data.jsonl` with this content (copy-paste into any text editor, save as `test_data.jsonl`):
|
||||
|
||||
```jsonl
|
||||
{"instruction": "What is Python?", "input": "", "output": "Python is a high-level programming language known for its simplicity."}
|
||||
{"instruction": "Explain recursion", "input": "", "output": "Recursion is when a function calls itself to solve smaller subproblems."}
|
||||
{"instruction": "What is a list?", "input": "", "output": "A list is an ordered, mutable collection of elements in Python."}
|
||||
{"instruction": "What is a dictionary?", "input": "", "output": "A dictionary is a key-value data structure in Python."}
|
||||
{"instruction": "What is OOP?", "input": "", "output": "OOP is a programming paradigm based on objects and classes."}
|
||||
{"instruction": "What is an API?", "input": "", "output": "An API is an interface that allows software systems to communicate."}
|
||||
{"instruction": "What is Git?", "input": "", "output": "Git is a distributed version control system for tracking code changes."}
|
||||
{"instruction": "What is Docker?", "input": "", "output": "Docker is a platform for containerizing applications."}
|
||||
{"instruction": "What is SQL?", "input": "", "output": "SQL is a language for managing and querying relational databases."}
|
||||
{"instruction": "What is REST?", "input": "", "output": "REST is an architectural style for designing networked APIs using HTTP methods."}
|
||||
```
|
||||
|
||||
Or create it with Python one-liner:
|
||||
|
||||
```cmd
|
||||
python -c "import json; data=[{'instruction':q,'input':'','output':a} for q,a in [('What is Python?','A high-level programming language.'),('Explain recursion','A function calling itself.'),('What is a list?','An ordered mutable collection.'),('What is OOP?','Programming with objects and classes.'),('What is Git?','A version control system.'),('What is Docker?','A containerization platform.'),('What is SQL?','A database query language.'),('What is REST?','An API architectural style.'),('What is an API?','An interface for software communication.'),('What is CSS?','A stylesheet language for web pages.')]]; f=open('test_data.jsonl','w'); [f.write(json.dumps(d)+'\n') for d in data]; f.close(); print('Created test_data.jsonl')"
|
||||
```
|
||||
|
||||
## 4. Data Tools
|
||||
|
||||
```cmd
|
||||
soup data inspect test_data.jsonl
|
||||
|
||||
soup data validate test_data.jsonl --format alpaca
|
||||
|
||||
soup data stats test_data.jsonl
|
||||
|
||||
soup data convert test_data.jsonl --to sharegpt -o test_sharegpt.jsonl
|
||||
|
||||
soup data convert test_data.jsonl --to chatml -o test_chatml.jsonl
|
||||
|
||||
soup data inspect test_sharegpt.jsonl
|
||||
|
||||
soup data inspect test_chatml.jsonl
|
||||
|
||||
soup data merge test_data.jsonl test_sharegpt.jsonl -o test_merged.jsonl --shuffle
|
||||
|
||||
soup data dedup test_merged.jsonl -o test_deduped.jsonl --threshold 0.8
|
||||
```
|
||||
|
||||
## 5. Create Config for Training
|
||||
|
||||
Create file `test_soup.yaml` (copy-paste into text editor):
|
||||
|
||||
```yaml
|
||||
base: TinyLlama/TinyLlama-1.1B-Chat-v1.0
|
||||
|
||||
data:
|
||||
train: test_data.jsonl
|
||||
format: alpaca
|
||||
max_length: 256
|
||||
|
||||
training:
|
||||
epochs: 2
|
||||
lr: 2e-4
|
||||
batch_size: 2
|
||||
quantization: 4bit
|
||||
logging_steps: 1
|
||||
save_steps: 50
|
||||
lora:
|
||||
r: 8
|
||||
alpha: 16
|
||||
dropout: 0.05
|
||||
|
||||
output: ./test_output
|
||||
```
|
||||
|
||||
## 6. Dry Run (validate without training)
|
||||
|
||||
```cmd
|
||||
soup train -c test_soup.yaml --dry-run
|
||||
```
|
||||
|
||||
## 7. Train
|
||||
|
||||
```cmd
|
||||
soup train -c test_soup.yaml --name "local-test"
|
||||
```
|
||||
|
||||
Training should take ~1-3 minutes on 3050 with this tiny dataset.
|
||||
|
||||
## 8. Experiment Tracking
|
||||
|
||||
```cmd
|
||||
soup runs
|
||||
```
|
||||
|
||||
Copy the Run ID from the output, then:
|
||||
|
||||
```cmd
|
||||
soup runs show RUN_ID_HERE
|
||||
```
|
||||
|
||||
Example: `soup runs show run_20260304_004948_983f284d`
|
||||
|
||||
## 9. Chat with Fine-Tuned Model
|
||||
|
||||
```cmd
|
||||
soup chat -m ./test_output
|
||||
```
|
||||
|
||||
Type questions, then type `exit` to quit.
|
||||
|
||||
## 10. Merge LoRA
|
||||
|
||||
```cmd
|
||||
soup merge -a ./test_output -o ./test_merged_model
|
||||
```
|
||||
|
||||
## 11. Export to GGUF (optional, needs llama.cpp + cmake)
|
||||
|
||||
```cmd
|
||||
soup export -m ./test_merged_model -q q4_k_m -o test_model.gguf
|
||||
```
|
||||
|
||||
## 12. Eval (optional, slow)
|
||||
|
||||
```cmd
|
||||
pip install lm-eval
|
||||
soup eval -m ./test_output --benchmarks hellaswag --batch-size 4
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cleanup (Windows)
|
||||
|
||||
```cmd
|
||||
rmdir /s /q test_output test_merged_model
|
||||
del test_data.jsonl test_sharegpt.jsonl test_chatml.jsonl test_merged.jsonl test_deduped.jsonl
|
||||
del test_soup.yaml test_chat.yaml test_code.yaml test_model.gguf
|
||||
```
|
||||
|
||||
Or in PowerShell:
|
||||
|
||||
```powershell
|
||||
Remove-Item -Recurse -Force test_output, test_merged_model -ErrorAction SilentlyContinue
|
||||
Remove-Item test_data.jsonl, test_sharegpt.jsonl, test_chatml.jsonl, test_merged.jsonl, test_deduped.jsonl, test_soup.yaml, test_chat.yaml, test_code.yaml, test_model.gguf -ErrorAction SilentlyContinue
|
||||
```
|
||||
|
||||
## Expected Results
|
||||
|
||||
| Step | Expected |
|
||||
|------|----------|
|
||||
| Version | `soup v0.2.1` |
|
||||
| Init templates | Creates yaml files |
|
||||
| Data inspect | Table with stats + sample rows |
|
||||
| Data validate | "20/20 rows valid" |
|
||||
| Data stats | Length distribution + histogram |
|
||||
| Data convert | Creates sharegpt/chatml jsonl files |
|
||||
| Data merge | Merges into single file |
|
||||
| Data dedup | Removes near-duplicates |
|
||||
| Dry run | "Config valid" or similar |
|
||||
| Train | Loss decreasing, ~1-3 min |
|
||||
| Runs | Shows run with metrics |
|
||||
| Chat | Model responds (quality low with 10 samples — that's OK) |
|
||||
| Merge | Creates full model in test_merged_model/ |
|
||||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.2.2"
|
||||
version = "0.3.0"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
|
@ -41,8 +41,11 @@ dependencies = [
|
|||
eval = ["lm-eval>=0.4.0"]
|
||||
data = ["datasketch>=1.6.0"]
|
||||
wandb = ["wandb>=0.15.0"]
|
||||
dev = ["pytest>=7.0", "ruff>=0.1.0", "pytest-cov>=4.0"]
|
||||
dev = ["pytest>=7.0", "ruff>=0.1.0", "pytest-cov>=4.0", "httpx>=0.24.0"]
|
||||
ui = ["fastapi>=0.104.0", "uvicorn>=0.24.0"]
|
||||
serve = ["fastapi>=0.104.0", "uvicorn>=0.24.0"]
|
||||
generate = ["httpx>=0.24.0"]
|
||||
deepspeed = ["deepspeed>=0.12.0"]
|
||||
|
||||
[project.scripts]
|
||||
soup = "soup_cli.cli:app"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.2.2"
|
||||
__version__ = "0.3.0"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,21 @@ import typer
|
|||
from rich.console import Console
|
||||
|
||||
from soup_cli import __version__
|
||||
from soup_cli.commands import chat, data, eval, export, init, merge, push, runs, train
|
||||
from soup_cli.commands import (
|
||||
chat,
|
||||
data,
|
||||
diff,
|
||||
eval,
|
||||
export,
|
||||
generate,
|
||||
init,
|
||||
merge,
|
||||
push,
|
||||
runs,
|
||||
serve,
|
||||
sweep,
|
||||
train,
|
||||
)
|
||||
|
||||
console = Console()
|
||||
|
||||
|
|
@ -28,6 +42,12 @@ app.add_typer(
|
|||
)
|
||||
app.add_typer(runs.app, name="runs", help="Experiment tracking: list, show, compare runs.")
|
||||
app.command(name="eval")(eval.eval_model)
|
||||
app.command()(serve.serve)
|
||||
app.command()(sweep.sweep)
|
||||
app.command(name="diff")(diff.diff)
|
||||
|
||||
# Register data generate as a subcommand of data
|
||||
data.app.command(name="generate")(generate.generate)
|
||||
|
||||
|
||||
@app.command()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,328 @@
|
|||
"""soup diff — compare outputs of two models side-by-side."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
from rich.columns import Columns
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def diff(
|
||||
model_a: str = typer.Option(
|
||||
...,
|
||||
"--model-a",
|
||||
"-a",
|
||||
help="Path to first model (LoRA adapter or full model)",
|
||||
),
|
||||
model_b: str = typer.Option(
|
||||
...,
|
||||
"--model-b",
|
||||
"-b",
|
||||
help="Path to second model (LoRA adapter or full model)",
|
||||
),
|
||||
prompts: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--prompts",
|
||||
"-p",
|
||||
help="Path to prompts file (JSONL with 'prompt' field, or one prompt per line)",
|
||||
),
|
||||
prompt: Optional[list[str]] = typer.Option(
|
||||
None,
|
||||
"--prompt",
|
||||
help="Single prompt to compare (can be repeated)",
|
||||
),
|
||||
base_a: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--base-a",
|
||||
help="Base model for model A (auto-detected for LoRA adapters)",
|
||||
),
|
||||
base_b: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--base-b",
|
||||
help="Base model for model B (auto-detected for LoRA adapters)",
|
||||
),
|
||||
max_tokens: int = typer.Option(
|
||||
256,
|
||||
"--max-tokens",
|
||||
help="Maximum tokens to generate per response",
|
||||
),
|
||||
temperature: float = typer.Option(
|
||||
0.7,
|
||||
"--temperature",
|
||||
"-t",
|
||||
help="Sampling temperature",
|
||||
),
|
||||
device: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--device",
|
||||
help="Device: cuda, mps, cpu. Auto-detected if not set.",
|
||||
),
|
||||
output: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--output",
|
||||
"-o",
|
||||
help="Save results to JSONL file",
|
||||
),
|
||||
):
|
||||
"""Compare outputs of two models side-by-side on the same prompts."""
|
||||
# Validate model paths
|
||||
path_a = Path(model_a)
|
||||
path_b = Path(model_b)
|
||||
if not path_a.exists():
|
||||
console.print(f"[red]Model A not found: {path_a}[/]")
|
||||
raise typer.Exit(1)
|
||||
if not path_b.exists():
|
||||
console.print(f"[red]Model B not found: {path_b}[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Collect prompts
|
||||
prompt_list = _collect_prompts(prompts, prompt)
|
||||
if not prompt_list:
|
||||
console.print("[red]No prompts provided. Use --prompts or --prompt.[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Detect device
|
||||
if not device:
|
||||
from soup_cli.utils.gpu import detect_device
|
||||
|
||||
device, _ = detect_device()
|
||||
|
||||
console.print(
|
||||
Panel(
|
||||
f"Model A: [bold]{path_a}[/]\n"
|
||||
f"Model B: [bold]{path_b}[/]\n"
|
||||
f"Prompts: [bold]{len(prompt_list)}[/]\n"
|
||||
f"Device: [bold]{device}[/]",
|
||||
title="Diff Plan",
|
||||
)
|
||||
)
|
||||
|
||||
# Load models
|
||||
console.print("[dim]Loading Model A...[/]")
|
||||
model_obj_a, tokenizer_a = _load_model(str(path_a), base_a, device)
|
||||
console.print("[dim]Loading Model B...[/]")
|
||||
model_obj_b, tokenizer_b = _load_model(str(path_b), base_b, device)
|
||||
console.print("[green]Both models loaded.[/]\n")
|
||||
|
||||
# Run comparison
|
||||
results = []
|
||||
for idx, prompt_text in enumerate(prompt_list):
|
||||
console.print(f"[bold]--- Prompt {idx + 1}/{len(prompt_list)} ---[/]")
|
||||
console.print(f"[dim]{prompt_text}[/]\n")
|
||||
|
||||
messages = [{"role": "user", "content": prompt_text}]
|
||||
|
||||
response_a = _generate(
|
||||
model_obj_a, tokenizer_a, messages,
|
||||
max_tokens=max_tokens, temperature=temperature,
|
||||
)
|
||||
response_b = _generate(
|
||||
model_obj_b, tokenizer_b, messages,
|
||||
max_tokens=max_tokens, temperature=temperature,
|
||||
)
|
||||
|
||||
# Side-by-side display
|
||||
panel_a = Panel(
|
||||
response_a or "[dim]<empty>[/]",
|
||||
title=f"[blue]Model A: {path_a.name}[/]",
|
||||
border_style="blue",
|
||||
width=console.width // 2 - 1,
|
||||
)
|
||||
panel_b = Panel(
|
||||
response_b or "[dim]<empty>[/]",
|
||||
title=f"[green]Model B: {path_b.name}[/]",
|
||||
border_style="green",
|
||||
width=console.width // 2 - 1,
|
||||
)
|
||||
console.print(Columns([panel_a, panel_b]))
|
||||
|
||||
# Metrics
|
||||
metrics = _compute_metrics(response_a, response_b)
|
||||
metrics_str = (
|
||||
f"Length: A={metrics['len_a']} / B={metrics['len_b']} | "
|
||||
f"Words: A={metrics['words_a']} / B={metrics['words_b']} | "
|
||||
f"Overlap: {metrics['word_overlap']:.0%}"
|
||||
)
|
||||
console.print(f"[dim]{metrics_str}[/]\n")
|
||||
|
||||
results.append({
|
||||
"prompt": prompt_text,
|
||||
"response_a": response_a,
|
||||
"response_b": response_b,
|
||||
"metrics": metrics,
|
||||
})
|
||||
|
||||
# Summary
|
||||
_display_summary(results, path_a.name, path_b.name)
|
||||
|
||||
# Save results
|
||||
if output:
|
||||
out_path = Path(output)
|
||||
with open(out_path, "w", encoding="utf-8") as f:
|
||||
for row in results:
|
||||
f.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
console.print(f"[dim]Results saved to {out_path}[/]")
|
||||
|
||||
|
||||
def _collect_prompts(prompts_file: Optional[str], prompt_args: Optional[list[str]]) -> list[str]:
|
||||
"""Collect prompts from file and/or CLI arguments."""
|
||||
result = []
|
||||
|
||||
if prompts_file:
|
||||
path = Path(prompts_file)
|
||||
if path.exists():
|
||||
with open(path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
# Try JSONL
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
if isinstance(obj, dict) and "prompt" in obj:
|
||||
result.append(obj["prompt"])
|
||||
continue
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
# Plain text
|
||||
result.append(line)
|
||||
|
||||
if prompt_args:
|
||||
result.extend(prompt_args)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _load_model(model_path: str, base_model: Optional[str], device: str):
|
||||
"""Load a model and tokenizer."""
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
path = Path(model_path)
|
||||
adapter_config_path = path / "adapter_config.json"
|
||||
is_adapter = adapter_config_path.exists()
|
||||
|
||||
if is_adapter and not base_model:
|
||||
try:
|
||||
with open(adapter_config_path, encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
base_model = config.get("base_model_name_or_path")
|
||||
except (json.JSONDecodeError, OSError):
|
||||
pass
|
||||
|
||||
if is_adapter and not base_model:
|
||||
console.print(f"[red]Cannot detect base model for {path}. Use --base-a/--base-b.[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
if is_adapter:
|
||||
from peft import PeftModel
|
||||
|
||||
base = AutoModelForCausalLM.from_pretrained(
|
||||
base_model,
|
||||
trust_remote_code=True,
|
||||
device_map="auto",
|
||||
dtype=torch.float16,
|
||||
)
|
||||
model_obj = PeftModel.from_pretrained(base, model_path)
|
||||
else:
|
||||
model_obj = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
device_map="auto",
|
||||
dtype=torch.float16,
|
||||
)
|
||||
|
||||
model_obj.eval()
|
||||
return model_obj, tokenizer
|
||||
|
||||
|
||||
def _generate(model, tokenizer, messages, max_tokens=256, temperature=0.7) -> str:
|
||||
"""Generate a response from the model."""
|
||||
import torch
|
||||
|
||||
if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
|
||||
text = tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
else:
|
||||
parts = []
|
||||
for msg in messages:
|
||||
role = msg["role"]
|
||||
content = msg["content"]
|
||||
if role == "system":
|
||||
parts.append(f"System: {content}")
|
||||
elif role == "user":
|
||||
parts.append(f"User: {content}")
|
||||
elif role == "assistant":
|
||||
parts.append(f"Assistant: {content}")
|
||||
parts.append("Assistant:")
|
||||
text = "\n".join(parts)
|
||||
|
||||
inputs = tokenizer(text, return_tensors="pt")
|
||||
input_ids = inputs["input_ids"].to(model.device)
|
||||
attention_mask = inputs["attention_mask"].to(model.device)
|
||||
|
||||
with torch.no_grad():
|
||||
gen_kwargs = {
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": attention_mask,
|
||||
"max_new_tokens": max_tokens,
|
||||
"do_sample": temperature > 0,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
}
|
||||
if temperature > 0:
|
||||
gen_kwargs["temperature"] = temperature
|
||||
gen_kwargs["top_p"] = 0.9
|
||||
outputs = model.generate(**gen_kwargs)
|
||||
|
||||
new_tokens = outputs[0][input_ids.shape[1]:]
|
||||
return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
|
||||
|
||||
|
||||
def _compute_metrics(response_a: str, response_b: str) -> dict:
|
||||
"""Compute comparison metrics between two responses."""
|
||||
words_a = set(response_a.lower().split())
|
||||
words_b = set(response_b.lower().split())
|
||||
|
||||
overlap = len(words_a & words_b) / max(len(words_a | words_b), 1)
|
||||
|
||||
return {
|
||||
"len_a": len(response_a),
|
||||
"len_b": len(response_b),
|
||||
"words_a": len(response_a.split()),
|
||||
"words_b": len(response_b.split()),
|
||||
"word_overlap": overlap,
|
||||
}
|
||||
|
||||
|
||||
def _display_summary(results: list[dict], name_a: str, name_b: str):
|
||||
"""Display summary statistics for the comparison."""
|
||||
if not results:
|
||||
return
|
||||
|
||||
table = Table(title="Comparison Summary")
|
||||
table.add_column("Metric", style="bold")
|
||||
table.add_column(f"Model A ({name_a})", justify="right", style="blue")
|
||||
table.add_column(f"Model B ({name_b})", justify="right", style="green")
|
||||
|
||||
avg_len_a = sum(r["metrics"]["len_a"] for r in results) / len(results)
|
||||
avg_len_b = sum(r["metrics"]["len_b"] for r in results) / len(results)
|
||||
avg_words_a = sum(r["metrics"]["words_a"] for r in results) / len(results)
|
||||
avg_words_b = sum(r["metrics"]["words_b"] for r in results) / len(results)
|
||||
avg_overlap = sum(r["metrics"]["word_overlap"] for r in results) / len(results)
|
||||
|
||||
table.add_row("Avg length (chars)", f"{avg_len_a:.0f}", f"{avg_len_b:.0f}")
|
||||
table.add_row("Avg words", f"{avg_words_a:.0f}", f"{avg_words_b:.0f}")
|
||||
table.add_row("Avg word overlap", f"{avg_overlap:.0%}", f"{avg_overlap:.0%}")
|
||||
|
||||
console.print(table)
|
||||
|
|
@ -0,0 +1,416 @@
|
|||
"""soup data generate — generate synthetic training data using LLMs."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def generate(
|
||||
prompt: str = typer.Option(
|
||||
...,
|
||||
"--prompt",
|
||||
"-p",
|
||||
help="System prompt describing what kind of data to generate",
|
||||
),
|
||||
count: int = typer.Option(
|
||||
100,
|
||||
"--count",
|
||||
"-n",
|
||||
help="Number of examples to generate",
|
||||
),
|
||||
output: str = typer.Option(
|
||||
"generated.jsonl",
|
||||
"--output",
|
||||
"-o",
|
||||
help="Output file path",
|
||||
),
|
||||
fmt: str = typer.Option(
|
||||
"alpaca",
|
||||
"--format",
|
||||
"-f",
|
||||
help="Output format: alpaca, sharegpt, chatml",
|
||||
),
|
||||
provider: str = typer.Option(
|
||||
"openai",
|
||||
"--provider",
|
||||
help="LLM provider: openai, local",
|
||||
),
|
||||
model_name: str = typer.Option(
|
||||
"gpt-4o-mini",
|
||||
"--model",
|
||||
"-m",
|
||||
help="Model name (OpenAI model ID or local model path)",
|
||||
),
|
||||
api_key: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--api-key",
|
||||
help="API key (or set OPENAI_API_KEY env var)",
|
||||
),
|
||||
api_base: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--api-base",
|
||||
help="Custom API base URL (for compatible APIs)",
|
||||
),
|
||||
batch_size: int = typer.Option(
|
||||
5,
|
||||
"--batch-size",
|
||||
help="Number of examples per API call",
|
||||
),
|
||||
temperature: float = typer.Option(
|
||||
0.8,
|
||||
"--temperature",
|
||||
"-t",
|
||||
help="Sampling temperature for generation",
|
||||
),
|
||||
dedup_with: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--dedup-with",
|
||||
help="Path to existing dataset to deduplicate against",
|
||||
),
|
||||
seed_file: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--seed",
|
||||
help="Path to seed examples file (JSONL) to guide generation",
|
||||
),
|
||||
):
|
||||
"""Generate synthetic training data using an LLM."""
|
||||
valid_formats = ("alpaca", "sharegpt", "chatml")
|
||||
if fmt not in valid_formats:
|
||||
console.print(f"[red]Invalid format: {fmt}. Must be one of: {', '.join(valid_formats)}[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
valid_providers = ("openai", "local")
|
||||
if provider not in valid_providers:
|
||||
console.print(
|
||||
f"[red]Invalid provider: {provider}. Must be one of: {', '.join(valid_providers)}[/]"
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Load seed examples if provided
|
||||
seed_examples = []
|
||||
if seed_file:
|
||||
seed_path = Path(seed_file)
|
||||
if not seed_path.exists():
|
||||
console.print(f"[red]Seed file not found: {seed_path}[/]")
|
||||
raise typer.Exit(1)
|
||||
from soup_cli.data.loader import load_raw_data
|
||||
|
||||
seed_examples = load_raw_data(seed_path)
|
||||
console.print(f"[dim]Loaded {len(seed_examples)} seed examples[/]")
|
||||
|
||||
# Load existing data for dedup
|
||||
existing_texts = set()
|
||||
if dedup_with:
|
||||
dedup_path = Path(dedup_with)
|
||||
if not dedup_path.exists():
|
||||
console.print(f"[red]Dedup file not found: {dedup_path}[/]")
|
||||
raise typer.Exit(1)
|
||||
from soup_cli.data.loader import load_raw_data
|
||||
|
||||
existing_data = load_raw_data(dedup_path)
|
||||
for row in existing_data:
|
||||
existing_texts.add(_row_to_text(row))
|
||||
console.print(f"[dim]Loaded {len(existing_texts)} existing examples for dedup[/]")
|
||||
|
||||
# Generate
|
||||
console.print(f"[dim]Generating {count} examples using {provider}/{model_name}...[/]")
|
||||
|
||||
all_examples = []
|
||||
duplicates = 0
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
task = progress.add_task("Generating...", total=count)
|
||||
|
||||
remaining = count
|
||||
while remaining > 0:
|
||||
current_batch = min(batch_size, remaining)
|
||||
|
||||
try:
|
||||
batch = _generate_batch(
|
||||
prompt=prompt,
|
||||
count=current_batch,
|
||||
fmt=fmt,
|
||||
provider=provider,
|
||||
model_name=model_name,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
temperature=temperature,
|
||||
seed_examples=seed_examples,
|
||||
)
|
||||
except Exception as exc:
|
||||
console.print(f"[red]Generation error: {exc}[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Validate and dedup
|
||||
for example in batch:
|
||||
if not _validate_example(example, fmt):
|
||||
continue
|
||||
text = _row_to_text(example)
|
||||
if text in existing_texts:
|
||||
duplicates += 1
|
||||
continue
|
||||
existing_texts.add(text)
|
||||
all_examples.append(example)
|
||||
|
||||
generated_this_round = len(batch)
|
||||
remaining -= current_batch
|
||||
progress.update(task, advance=generated_this_round)
|
||||
|
||||
# Write output
|
||||
out_path = Path(output)
|
||||
with open(out_path, "w", encoding="utf-8") as f:
|
||||
for row in all_examples:
|
||||
f.write(json.dumps(row, ensure_ascii=False) + "\n")
|
||||
|
||||
console.print(
|
||||
f"\n[green]Generated {len(all_examples)} examples[/]\n"
|
||||
f"Format: [bold]{fmt}[/]\n"
|
||||
f"Output: [bold]{out_path}[/]\n"
|
||||
+ (f"Duplicates: [yellow]{duplicates} removed[/]\n" if duplicates > 0 else "")
|
||||
)
|
||||
|
||||
|
||||
def _generate_batch(
|
||||
prompt: str,
|
||||
count: int,
|
||||
fmt: str,
|
||||
provider: str,
|
||||
model_name: str,
|
||||
api_key: Optional[str],
|
||||
api_base: Optional[str],
|
||||
temperature: float,
|
||||
seed_examples: list[dict],
|
||||
) -> list[dict]:
|
||||
"""Generate a batch of examples using the specified provider."""
|
||||
if provider == "openai":
|
||||
return _generate_openai(
|
||||
prompt=prompt,
|
||||
count=count,
|
||||
fmt=fmt,
|
||||
model_name=model_name,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
temperature=temperature,
|
||||
seed_examples=seed_examples,
|
||||
)
|
||||
elif provider == "local":
|
||||
return _generate_local(
|
||||
prompt=prompt,
|
||||
count=count,
|
||||
fmt=fmt,
|
||||
model_name=model_name,
|
||||
temperature=temperature,
|
||||
seed_examples=seed_examples,
|
||||
)
|
||||
return []
|
||||
|
||||
|
||||
def _build_generation_prompt(prompt: str, count: int, fmt: str, seed_examples: list) -> str:
|
||||
"""Build the prompt for data generation."""
|
||||
format_spec = {
|
||||
"alpaca": (
|
||||
'Each example must be a JSON object with keys: '
|
||||
'"instruction", "input" (can be empty string), "output".'
|
||||
),
|
||||
"sharegpt": (
|
||||
'Each example must be a JSON object with key "conversations", '
|
||||
'which is a list of objects with "from" (human/gpt) and "value".'
|
||||
),
|
||||
"chatml": (
|
||||
'Each example must be a JSON object with key "messages", '
|
||||
'which is a list of objects with "role" (user/assistant) and "content".'
|
||||
),
|
||||
}
|
||||
|
||||
system_msg = (
|
||||
f"You are a training data generator. Generate exactly {count} diverse, "
|
||||
f"high-quality training examples.\n\n"
|
||||
f"Topic/Instructions: {prompt}\n\n"
|
||||
f"Format: {format_spec[fmt]}\n\n"
|
||||
f"Return ONLY a JSON array of {count} examples. No markdown, no explanation."
|
||||
)
|
||||
|
||||
if seed_examples:
|
||||
seed_str = json.dumps(seed_examples[:3], ensure_ascii=False, indent=2)
|
||||
system_msg += f"\n\nHere are some seed examples to guide the style:\n{seed_str}"
|
||||
|
||||
return system_msg
|
||||
|
||||
|
||||
def _generate_openai(
|
||||
prompt: str,
|
||||
count: int,
|
||||
fmt: str,
|
||||
model_name: str,
|
||||
api_key: Optional[str],
|
||||
api_base: Optional[str],
|
||||
temperature: float,
|
||||
seed_examples: list[dict],
|
||||
) -> list[dict]:
|
||||
"""Generate examples using OpenAI-compatible API."""
|
||||
import os
|
||||
|
||||
resolved_key = api_key or os.environ.get("OPENAI_API_KEY")
|
||||
if not resolved_key:
|
||||
raise ValueError(
|
||||
"OpenAI API key not found. Set OPENAI_API_KEY env var or pass --api-key."
|
||||
)
|
||||
|
||||
try:
|
||||
import httpx
|
||||
except ImportError:
|
||||
raise ImportError("httpx is required for OpenAI generation. Install: pip install httpx")
|
||||
|
||||
base_url = api_base or "https://api.openai.com/v1"
|
||||
generation_prompt = _build_generation_prompt(prompt, count, fmt, seed_examples)
|
||||
|
||||
response = httpx.post(
|
||||
f"{base_url}/chat/completions",
|
||||
headers={
|
||||
"Authorization": f"Bearer {resolved_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": [
|
||||
{"role": "system", "content": generation_prompt},
|
||||
{"role": "user", "content": f"Generate {count} training examples now."},
|
||||
],
|
||||
"temperature": temperature,
|
||||
"max_tokens": 4096,
|
||||
},
|
||||
timeout=120.0,
|
||||
)
|
||||
|
||||
if response.status_code != 200:
|
||||
raise ValueError(f"API returned {response.status_code}: {response.text}")
|
||||
|
||||
data = response.json()
|
||||
content = data["choices"][0]["message"]["content"]
|
||||
|
||||
return _parse_json_array(content)
|
||||
|
||||
|
||||
def _generate_local(
|
||||
prompt: str,
|
||||
count: int,
|
||||
fmt: str,
|
||||
model_name: str,
|
||||
temperature: float,
|
||||
seed_examples: list[dict],
|
||||
) -> list[dict]:
|
||||
"""Generate examples using a local model via transformers."""
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
trust_remote_code=True,
|
||||
device_map="auto",
|
||||
torch_dtype=torch.float16,
|
||||
)
|
||||
model.eval()
|
||||
|
||||
generation_prompt = _build_generation_prompt(prompt, count, fmt, seed_examples)
|
||||
|
||||
if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
|
||||
messages = [
|
||||
{"role": "system", "content": generation_prompt},
|
||||
{"role": "user", "content": f"Generate {count} training examples now."},
|
||||
]
|
||||
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||
else:
|
||||
text = f"{generation_prompt}\n\nGenerate {count} training examples now.\n\n"
|
||||
|
||||
inputs = tokenizer(text, return_tensors="pt")
|
||||
input_ids = inputs["input_ids"].to(model.device)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model.generate(
|
||||
input_ids,
|
||||
max_new_tokens=4096,
|
||||
do_sample=temperature > 0,
|
||||
temperature=temperature if temperature > 0 else None,
|
||||
top_p=0.9 if temperature > 0 else None,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
|
||||
new_tokens = outputs[0][input_ids.shape[1]:]
|
||||
content = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
|
||||
|
||||
return _parse_json_array(content)
|
||||
|
||||
|
||||
def _parse_json_array(content: str) -> list[dict]:
|
||||
"""Parse a JSON array from LLM output, handling markdown code blocks."""
|
||||
content = content.strip()
|
||||
|
||||
# Strip markdown code fences
|
||||
if content.startswith("```"):
|
||||
lines = content.split("\n")
|
||||
# Remove first line (```json or ```)
|
||||
lines = lines[1:]
|
||||
# Remove last line if it's ```)
|
||||
if lines and lines[-1].strip() == "```":
|
||||
lines = lines[:-1]
|
||||
content = "\n".join(lines).strip()
|
||||
|
||||
# Try to find JSON array in content
|
||||
start = content.find("[")
|
||||
end = content.rfind("]")
|
||||
if start != -1 and end != -1 and end > start:
|
||||
content = content[start:end + 1]
|
||||
|
||||
try:
|
||||
result = json.loads(content)
|
||||
if isinstance(result, list):
|
||||
return [item for item in result if isinstance(item, dict)]
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Try line-by-line JSON objects
|
||||
results = []
|
||||
for line in content.split("\n"):
|
||||
line = line.strip()
|
||||
if line.startswith("{"):
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
if isinstance(obj, dict):
|
||||
results.append(obj)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
return results
|
||||
|
||||
|
||||
def _validate_example(example: dict, fmt: str) -> bool:
|
||||
"""Validate a single generated example matches the expected format."""
|
||||
if fmt == "alpaca":
|
||||
return "instruction" in example and "output" in example
|
||||
elif fmt == "sharegpt":
|
||||
convos = example.get("conversations", [])
|
||||
return len(convos) >= 2
|
||||
elif fmt == "chatml":
|
||||
msgs = example.get("messages", [])
|
||||
return len(msgs) >= 2
|
||||
return False
|
||||
|
||||
|
||||
def _row_to_text(row: dict) -> str:
|
||||
"""Convert a row to a text string for dedup comparison."""
|
||||
return " ".join(str(v) for v in row.values() if v)
|
||||
|
|
@ -0,0 +1,404 @@
|
|||
"""soup serve — local inference server with OpenAI-compatible API."""
|
||||
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def serve(
|
||||
model: str = typer.Option(
|
||||
...,
|
||||
"--model",
|
||||
"-m",
|
||||
help="Path to LoRA adapter directory or full model",
|
||||
),
|
||||
base_model: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--base",
|
||||
"-b",
|
||||
help="Base model ID. Auto-detected from adapter_config.json if not set.",
|
||||
),
|
||||
port: int = typer.Option(
|
||||
8000,
|
||||
"--port",
|
||||
"-p",
|
||||
help="Port to serve on",
|
||||
),
|
||||
host: str = typer.Option(
|
||||
"0.0.0.0",
|
||||
"--host",
|
||||
help="Host to bind to",
|
||||
),
|
||||
device: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--device",
|
||||
help="Device: cuda, mps, cpu. Auto-detected if not set.",
|
||||
),
|
||||
max_tokens_default: int = typer.Option(
|
||||
512,
|
||||
"--max-tokens",
|
||||
help="Default max tokens for generation",
|
||||
),
|
||||
):
|
||||
"""Start a local inference server with OpenAI-compatible API."""
|
||||
# Lazy imports for fast CLI startup
|
||||
try:
|
||||
import uvicorn # noqa: F401
|
||||
from fastapi import FastAPI # noqa: F401
|
||||
from fastapi.responses import StreamingResponse # noqa: F401
|
||||
except ImportError:
|
||||
console.print(
|
||||
"[red]FastAPI/uvicorn not installed.[/]\n"
|
||||
"Install with: [bold]pip install 'soup-cli[ui]'[/]"
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
model_path = Path(model)
|
||||
if not model_path.exists():
|
||||
console.print(f"[red]Model path not found: {model_path}[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Detect adapter
|
||||
adapter_config_path = model_path / "adapter_config.json"
|
||||
is_adapter = adapter_config_path.exists()
|
||||
|
||||
# Resolve base model
|
||||
if is_adapter and not base_model:
|
||||
base_model = _detect_base_model(adapter_config_path)
|
||||
if not base_model:
|
||||
console.print(
|
||||
"[red]Cannot detect base model from adapter_config.json.[/]\n"
|
||||
"Please specify with [bold]--base[/] flag."
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Detect device
|
||||
if not device:
|
||||
from soup_cli.utils.gpu import detect_device
|
||||
|
||||
device, _ = detect_device()
|
||||
|
||||
console.print(
|
||||
Panel(
|
||||
f"Model: [bold]{model_path}[/]\n"
|
||||
+ (f"Base: [bold]{base_model}[/]\n" if is_adapter else "")
|
||||
+ f"Device: [bold]{device}[/]\n"
|
||||
f"Type: [bold]{'LoRA adapter' if is_adapter else 'Full model'}[/]",
|
||||
title="Loading model",
|
||||
)
|
||||
)
|
||||
|
||||
# Load model
|
||||
model_obj, tokenizer = _load_model(
|
||||
model_path=str(model_path),
|
||||
base_model=base_model,
|
||||
is_adapter=is_adapter,
|
||||
device=device,
|
||||
)
|
||||
console.print("[bold green]Model loaded![/]")
|
||||
|
||||
# Build and start FastAPI app
|
||||
app = _create_app(
|
||||
model_obj=model_obj,
|
||||
tokenizer=tokenizer,
|
||||
device=device,
|
||||
model_name=str(model_path.name),
|
||||
max_tokens_default=max_tokens_default,
|
||||
)
|
||||
|
||||
console.print(
|
||||
Panel(
|
||||
f"URL: [bold]http://{host}:{port}[/]\n"
|
||||
f"Endpoints: [bold]/v1/chat/completions[/], [bold]/v1/models[/], [bold]/health[/]\n\n"
|
||||
f"Example:\n"
|
||||
f" curl http://localhost:{port}/v1/chat/completions \\\n"
|
||||
f' -H "Content-Type: application/json" \\\n'
|
||||
f" -d '{{"
|
||||
f'"model": "{model_path.name}", '
|
||||
f'"messages": [{{"role": "user", "content": "Hello!"}}]'
|
||||
f"}}'\n\n"
|
||||
f"Press [bold]Ctrl+C[/] to stop.",
|
||||
title="[bold green]Server Ready[/]",
|
||||
)
|
||||
)
|
||||
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host=host, port=port, log_level="warning")
|
||||
|
||||
|
||||
def _detect_base_model(adapter_config_path: Path) -> Optional[str]:
|
||||
"""Read base_model_name_or_path from adapter_config.json."""
|
||||
try:
|
||||
with open(adapter_config_path, encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
return config.get("base_model_name_or_path")
|
||||
except (json.JSONDecodeError, OSError):
|
||||
return None
|
||||
|
||||
|
||||
def _load_model(
|
||||
model_path: str,
|
||||
base_model: Optional[str],
|
||||
is_adapter: bool,
|
||||
device: str,
|
||||
):
|
||||
"""Load model and tokenizer."""
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
console.print("[dim]Loading tokenizer...[/]")
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
if is_adapter:
|
||||
from peft import PeftModel
|
||||
|
||||
console.print(f"[dim]Loading base model: {base_model}...[/]")
|
||||
base = AutoModelForCausalLM.from_pretrained(
|
||||
base_model,
|
||||
trust_remote_code=True,
|
||||
device_map="auto",
|
||||
dtype=torch.float16,
|
||||
)
|
||||
console.print(f"[dim]Loading LoRA adapter: {model_path}...[/]")
|
||||
model_obj = PeftModel.from_pretrained(base, model_path)
|
||||
else:
|
||||
console.print(f"[dim]Loading model: {model_path}...[/]")
|
||||
model_obj = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
device_map="auto",
|
||||
dtype=torch.float16,
|
||||
)
|
||||
|
||||
model_obj.eval()
|
||||
return model_obj, tokenizer
|
||||
|
||||
|
||||
def _generate_response(
|
||||
model,
|
||||
tokenizer,
|
||||
messages: list[dict],
|
||||
max_tokens: int = 512,
|
||||
temperature: float = 0.7,
|
||||
top_p: float = 0.9,
|
||||
stream: bool = False,
|
||||
):
|
||||
"""Generate a response from the model."""
|
||||
import torch
|
||||
|
||||
# Apply chat template
|
||||
if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
|
||||
text = tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
else:
|
||||
parts = []
|
||||
for msg in messages:
|
||||
role = msg["role"]
|
||||
content = msg["content"]
|
||||
if role == "system":
|
||||
parts.append(f"System: {content}")
|
||||
elif role == "user":
|
||||
parts.append(f"User: {content}")
|
||||
elif role == "assistant":
|
||||
parts.append(f"Assistant: {content}")
|
||||
parts.append("Assistant:")
|
||||
text = "\n".join(parts)
|
||||
|
||||
inputs = tokenizer(text, return_tensors="pt")
|
||||
input_ids = inputs["input_ids"].to(model.device)
|
||||
attention_mask = inputs["attention_mask"].to(model.device)
|
||||
|
||||
with torch.no_grad():
|
||||
gen_kwargs = {
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": attention_mask,
|
||||
"max_new_tokens": max_tokens,
|
||||
"do_sample": temperature > 0,
|
||||
"pad_token_id": tokenizer.pad_token_id,
|
||||
}
|
||||
if temperature > 0:
|
||||
gen_kwargs["temperature"] = temperature
|
||||
gen_kwargs["top_p"] = top_p
|
||||
|
||||
outputs = model.generate(**gen_kwargs)
|
||||
|
||||
new_tokens = outputs[0][input_ids.shape[1]:]
|
||||
response = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
|
||||
|
||||
prompt_tokens = input_ids.shape[1]
|
||||
completion_tokens = len(new_tokens)
|
||||
|
||||
return response, prompt_tokens, completion_tokens
|
||||
|
||||
|
||||
def _create_app(
|
||||
model_obj,
|
||||
tokenizer,
|
||||
device: str,
|
||||
model_name: str,
|
||||
max_tokens_default: int,
|
||||
):
|
||||
"""Create the FastAPI application with OpenAI-compatible endpoints."""
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel as PydanticBaseModel
|
||||
from pydantic import Field
|
||||
|
||||
app = FastAPI(title="Soup Inference Server", version="1.0.0")
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# --- Request/Response models ---
|
||||
|
||||
class ChatMessage(PydanticBaseModel):
|
||||
role: str
|
||||
content: str
|
||||
|
||||
class ChatCompletionRequest(PydanticBaseModel):
|
||||
model: str = model_name
|
||||
messages: list[ChatMessage]
|
||||
temperature: float = Field(default=0.7, ge=0.0, le=2.0)
|
||||
top_p: float = Field(default=0.9, ge=0.0, le=1.0)
|
||||
max_tokens: Optional[int] = None
|
||||
stream: bool = False
|
||||
|
||||
# --- Endpoints ---
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"status": "ok", "model": model_name, "device": device}
|
||||
|
||||
@app.get("/v1/models")
|
||||
def list_models():
|
||||
return {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": model_name,
|
||||
"object": "model",
|
||||
"owned_by": "soup",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
def chat_completions(request: ChatCompletionRequest):
|
||||
messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
|
||||
max_tokens = request.max_tokens or max_tokens_default
|
||||
|
||||
if request.stream:
|
||||
return StreamingResponse(
|
||||
_stream_response(
|
||||
model_obj, tokenizer, messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=request.temperature,
|
||||
top_p=request.top_p,
|
||||
model_name=model_name,
|
||||
),
|
||||
media_type="text/event-stream",
|
||||
)
|
||||
|
||||
try:
|
||||
response_text, prompt_tokens, completion_tokens = _generate_response(
|
||||
model_obj, tokenizer, messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=request.temperature,
|
||||
top_p=request.top_p,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=str(exc))
|
||||
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": model_name,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": response_text,
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
},
|
||||
}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def _stream_response(
|
||||
model, tokenizer, messages,
|
||||
max_tokens, temperature, top_p, model_name,
|
||||
):
|
||||
"""Generator that yields SSE chunks for streaming responses."""
|
||||
chat_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
|
||||
created = int(time.time())
|
||||
|
||||
# Generate full response (true token-by-token streaming requires TextIteratorStreamer)
|
||||
response_text, _, _ = _generate_response(
|
||||
model, tokenizer, messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
)
|
||||
|
||||
# Simulate streaming by sending word-by-word
|
||||
words = response_text.split(" ")
|
||||
for idx, word in enumerate(words):
|
||||
chunk_text = word if idx == 0 else f" {word}"
|
||||
chunk = {
|
||||
"id": chat_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": chunk_text},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
yield f"data: {json.dumps(chunk)}\n\n"
|
||||
|
||||
# Final chunk
|
||||
final_chunk = {
|
||||
"id": chat_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created,
|
||||
"model": model_name,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
}
|
||||
yield f"data: {json.dumps(final_chunk)}\n\n"
|
||||
yield "data: [DONE]\n\n"
|
||||
|
|
@ -0,0 +1,360 @@
|
|||
"""soup sweep — hyperparameter search over training configs."""
|
||||
|
||||
import itertools
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
|
||||
from soup_cli.config.loader import load_config
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def sweep(
|
||||
config: str = typer.Option(
|
||||
"soup.yaml",
|
||||
"--config",
|
||||
"-c",
|
||||
help="Path to base soup.yaml config file",
|
||||
),
|
||||
param: list[str] = typer.Option(
|
||||
...,
|
||||
"--param",
|
||||
"-p",
|
||||
help="Parameter to sweep: key=val1,val2,val3 (e.g., lr=1e-5,2e-5,5e-5)",
|
||||
),
|
||||
strategy: str = typer.Option(
|
||||
"grid",
|
||||
"--strategy",
|
||||
"-s",
|
||||
help="Search strategy: grid, random",
|
||||
),
|
||||
max_runs: Optional[int] = typer.Option(
|
||||
None,
|
||||
"--max-runs",
|
||||
help="Max number of runs (useful for random strategy)",
|
||||
),
|
||||
name: Optional[str] = typer.Option(
|
||||
None,
|
||||
"--name",
|
||||
"-n",
|
||||
help="Sweep experiment name prefix",
|
||||
),
|
||||
dry_run: bool = typer.Option(
|
||||
False,
|
||||
"--dry-run",
|
||||
help="Show planned runs without executing",
|
||||
),
|
||||
):
|
||||
"""Run hyperparameter sweep: grid or random search over training parameters."""
|
||||
config_path = Path(config)
|
||||
if not config_path.exists():
|
||||
console.print(f"[red]Config not found: {config_path}[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
if strategy not in ("grid", "random"):
|
||||
console.print(f"[red]Invalid strategy: {strategy}. Must be grid or random.[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Parse sweep parameters
|
||||
sweep_params = _parse_sweep_params(param)
|
||||
if not sweep_params:
|
||||
console.print("[red]No valid sweep parameters provided.[/]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Generate parameter combinations
|
||||
combinations = _generate_combinations(sweep_params, strategy, max_runs)
|
||||
|
||||
console.print(
|
||||
Panel(
|
||||
f"Config: [bold]{config_path}[/]\n"
|
||||
f"Strategy: [bold]{strategy}[/]\n"
|
||||
f"Params: [bold]{', '.join(sweep_params.keys())}[/]\n"
|
||||
f"Runs: [bold]{len(combinations)}[/]",
|
||||
title="Sweep Plan",
|
||||
)
|
||||
)
|
||||
|
||||
# Show parameter table
|
||||
param_table = Table(title="Parameter Grid")
|
||||
param_table.add_column("Run", style="bold")
|
||||
for key in sweep_params:
|
||||
param_table.add_column(key)
|
||||
|
||||
for idx, combo in enumerate(combinations):
|
||||
row_values = [str(combo[key]) for key in sweep_params]
|
||||
param_table.add_row(f"#{idx + 1}", *row_values)
|
||||
|
||||
console.print(param_table)
|
||||
|
||||
if dry_run:
|
||||
console.print("[yellow]Dry run — no training will be executed.[/]")
|
||||
raise typer.Exit()
|
||||
|
||||
# Execute sweep
|
||||
base_cfg = load_config(config_path)
|
||||
results = []
|
||||
|
||||
for idx, combo in enumerate(combinations):
|
||||
run_name = f"{name or 'sweep'}_{idx + 1}"
|
||||
console.print(f"\n[bold]--- Run {idx + 1}/{len(combinations)}: {run_name} ---[/]")
|
||||
|
||||
for key, val in combo.items():
|
||||
console.print(f" {key} = {val}")
|
||||
|
||||
try:
|
||||
result = _run_single(base_cfg, combo, run_name, config_path)
|
||||
results.append({
|
||||
"name": run_name,
|
||||
"params": combo,
|
||||
"run_id": result.get("run_id", ""),
|
||||
"final_loss": result.get("final_loss", 0),
|
||||
"duration": result.get("duration", ""),
|
||||
"status": "completed",
|
||||
})
|
||||
except Exception as exc:
|
||||
console.print(f"[red]Run {run_name} failed: {exc}[/]")
|
||||
results.append({
|
||||
"name": run_name,
|
||||
"params": combo,
|
||||
"run_id": "",
|
||||
"final_loss": 0,
|
||||
"duration": "",
|
||||
"status": "failed",
|
||||
})
|
||||
|
||||
# Summary table
|
||||
_display_summary(results, sweep_params)
|
||||
|
||||
|
||||
def _parse_sweep_params(params: list[str]) -> dict[str, list]:
|
||||
"""Parse sweep parameter strings into a dict of {key: [values]}."""
|
||||
result = {}
|
||||
for param_str in params:
|
||||
if "=" not in param_str:
|
||||
console.print(f"[yellow]Skipping invalid param: {param_str} (missing '=')[/]")
|
||||
continue
|
||||
|
||||
key, values_str = param_str.split("=", 1)
|
||||
key = key.strip()
|
||||
values = []
|
||||
|
||||
for val in values_str.split(","):
|
||||
val = val.strip()
|
||||
values.append(_parse_value(val))
|
||||
|
||||
if values:
|
||||
result[key] = values
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _parse_value(val: str):
|
||||
"""Parse a string value into the appropriate Python type."""
|
||||
# Bool
|
||||
if val.lower() in ("true", "false"):
|
||||
return val.lower() == "true"
|
||||
# None
|
||||
if val.lower() == "none":
|
||||
return None
|
||||
# Int
|
||||
try:
|
||||
return int(val)
|
||||
except ValueError:
|
||||
pass
|
||||
# Float (including scientific notation)
|
||||
try:
|
||||
return float(val)
|
||||
except ValueError:
|
||||
pass
|
||||
# String
|
||||
return val
|
||||
|
||||
|
||||
def _generate_combinations(
|
||||
sweep_params: dict[str, list],
|
||||
strategy: str,
|
||||
max_runs: Optional[int],
|
||||
) -> list[dict]:
|
||||
"""Generate parameter combinations based on strategy."""
|
||||
keys = list(sweep_params.keys())
|
||||
value_lists = [sweep_params[k] for k in keys]
|
||||
|
||||
if strategy == "grid":
|
||||
combos = [dict(zip(keys, vals)) for vals in itertools.product(*value_lists)]
|
||||
elif strategy == "random":
|
||||
total_possible = 1
|
||||
for vals in value_lists:
|
||||
total_possible *= len(vals)
|
||||
|
||||
num_runs = max_runs or min(total_possible, 10)
|
||||
num_runs = min(num_runs, total_possible)
|
||||
|
||||
if num_runs >= total_possible:
|
||||
# Just do all of them
|
||||
combos = [dict(zip(keys, vals)) for vals in itertools.product(*value_lists)]
|
||||
else:
|
||||
seen = set()
|
||||
combos = []
|
||||
while len(combos) < num_runs:
|
||||
vals = tuple(random.choice(vals_list) for vals_list in value_lists)
|
||||
if vals not in seen:
|
||||
seen.add(vals)
|
||||
combos.append(dict(zip(keys, vals)))
|
||||
else:
|
||||
combos = []
|
||||
|
||||
if max_runs and len(combos) > max_runs:
|
||||
combos = combos[:max_runs]
|
||||
|
||||
return combos
|
||||
|
||||
|
||||
def _set_nested_param(config_dict: dict, key: str, value) -> dict:
|
||||
"""Set a nested parameter in a config dict using dot notation.
|
||||
|
||||
Supports keys like: lr, lora.r, training.epochs, etc.
|
||||
Maps common short names to their full paths.
|
||||
"""
|
||||
# Short name mappings
|
||||
shortcuts = {
|
||||
"lr": "training.lr",
|
||||
"epochs": "training.epochs",
|
||||
"batch_size": "training.batch_size",
|
||||
"lora_r": "training.lora.r",
|
||||
"lora_alpha": "training.lora.alpha",
|
||||
"lora_dropout": "training.lora.dropout",
|
||||
"quantization": "training.quantization",
|
||||
"warmup_ratio": "training.warmup_ratio",
|
||||
"weight_decay": "training.weight_decay",
|
||||
"gradient_accumulation_steps": "training.gradient_accumulation_steps",
|
||||
"max_grad_norm": "training.max_grad_norm",
|
||||
"optimizer": "training.optimizer",
|
||||
"scheduler": "training.scheduler",
|
||||
"val_split": "data.val_split",
|
||||
"max_length": "data.max_length",
|
||||
"dpo_beta": "training.dpo_beta",
|
||||
}
|
||||
|
||||
full_key = shortcuts.get(key, key)
|
||||
parts = full_key.split(".")
|
||||
|
||||
obj = config_dict
|
||||
for part in parts[:-1]:
|
||||
if part not in obj:
|
||||
obj[part] = {}
|
||||
obj = obj[part]
|
||||
obj[parts[-1]] = value
|
||||
|
||||
return config_dict
|
||||
|
||||
|
||||
def _run_single(base_cfg, params: dict, run_name: str, config_path: Path) -> dict:
|
||||
"""Run a single training with modified parameters."""
|
||||
from soup_cli.config.schema import SoupConfig
|
||||
from soup_cli.data.loader import load_dataset
|
||||
from soup_cli.experiment.tracker import ExperimentTracker
|
||||
from soup_cli.monitoring.display import TrainingDisplay
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
from soup_cli.utils.gpu import detect_device, get_gpu_info
|
||||
|
||||
# Deep copy and modify config
|
||||
config_dict = base_cfg.model_dump()
|
||||
for key, val in params.items():
|
||||
_set_nested_param(config_dict, key, val)
|
||||
|
||||
# Override experiment name
|
||||
config_dict["experiment_name"] = run_name
|
||||
cfg = SoupConfig(**config_dict)
|
||||
|
||||
# Detect hardware
|
||||
device, device_name = detect_device()
|
||||
gpu_info = get_gpu_info()
|
||||
|
||||
# Load data
|
||||
dataset = load_dataset(cfg.data)
|
||||
console.print(f"[dim]Loaded {len(dataset['train'])} train samples[/]")
|
||||
|
||||
# Start tracking
|
||||
tracker = ExperimentTracker()
|
||||
run_id = tracker.start_run(
|
||||
config_dict=cfg.model_dump(),
|
||||
device=device,
|
||||
device_name=device_name,
|
||||
gpu_info=gpu_info,
|
||||
experiment_name=run_name,
|
||||
)
|
||||
|
||||
# Build trainer
|
||||
if cfg.task == "dpo":
|
||||
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
||||
|
||||
trainer_wrapper = DPOTrainerWrapper(cfg, device=device)
|
||||
else:
|
||||
trainer_wrapper = SFTTrainerWrapper(cfg, device=device)
|
||||
trainer_wrapper.setup(dataset)
|
||||
|
||||
# Train
|
||||
display = TrainingDisplay(cfg, device_name=device_name)
|
||||
try:
|
||||
result = trainer_wrapper.train(display=display, tracker=tracker, run_id=run_id)
|
||||
tracker.finish_run(
|
||||
run_id=run_id,
|
||||
initial_loss=result["initial_loss"],
|
||||
final_loss=result["final_loss"],
|
||||
total_steps=result["total_steps"],
|
||||
duration_secs=result["duration_secs"],
|
||||
output_dir=result["output_dir"],
|
||||
)
|
||||
result["run_id"] = run_id
|
||||
return result
|
||||
except Exception:
|
||||
tracker.fail_run(run_id)
|
||||
raise
|
||||
|
||||
|
||||
def _display_summary(results: list[dict], sweep_params: dict[str, list]):
|
||||
"""Display sweep results summary table."""
|
||||
table = Table(title="Sweep Results")
|
||||
table.add_column("Run", style="bold")
|
||||
for key in sweep_params:
|
||||
table.add_column(key)
|
||||
table.add_column("Final Loss", justify="right", style="green")
|
||||
table.add_column("Duration", justify="right")
|
||||
table.add_column("Status")
|
||||
|
||||
# Sort by final loss (best first)
|
||||
sorted_results = sorted(results, key=lambda r: r.get("final_loss", float("inf")))
|
||||
|
||||
for idx, res in enumerate(sorted_results):
|
||||
status_style = "green" if res["status"] == "completed" else "red"
|
||||
param_vals = [str(res["params"].get(k, "")) for k in sweep_params]
|
||||
loss_str = f"{res['final_loss']:.4f}" if res["final_loss"] else "—"
|
||||
best_marker = " [bold yellow]*[/]" if idx == 0 and res["status"] == "completed" else ""
|
||||
table.add_row(
|
||||
res["name"],
|
||||
*param_vals,
|
||||
f"{loss_str}{best_marker}",
|
||||
res.get("duration", "—"),
|
||||
f"[{status_style}]{res['status']}[/]",
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
# Best run
|
||||
completed = [r for r in sorted_results if r["status"] == "completed"]
|
||||
if completed:
|
||||
best = completed[0]
|
||||
console.print(
|
||||
f"\n[bold green]Best run:[/] {best['name']} "
|
||||
f"(loss: {best['final_loss']:.4f})"
|
||||
)
|
||||
for key, val in best["params"].items():
|
||||
console.print(f" {key} = {val}")
|
||||
if best.get("run_id"):
|
||||
console.print(f"\n[dim]View details: soup runs show {best['run_id']}[/]")
|
||||
|
|
@ -44,6 +44,11 @@ def train(
|
|||
"--wandb",
|
||||
help="Enable Weights & Biases logging",
|
||||
),
|
||||
deepspeed: str = typer.Option(
|
||||
None,
|
||||
"--deepspeed",
|
||||
help="Enable DeepSpeed: zero2, zero3, zero2_offload, or path to config JSON",
|
||||
),
|
||||
):
|
||||
"""Start training from a soup.yaml config."""
|
||||
config_path = Path(config)
|
||||
|
|
@ -79,6 +84,13 @@ def train(
|
|||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
# --- DeepSpeed setup ---
|
||||
ds_config_path = None
|
||||
if deepspeed:
|
||||
ds_config_path = _resolve_deepspeed(deepspeed)
|
||||
if ds_config_path:
|
||||
console.print(f"[green]DeepSpeed enabled:[/] {deepspeed}")
|
||||
|
||||
# Detect hardware
|
||||
device, device_name = detect_device()
|
||||
gpu_info = get_gpu_info()
|
||||
|
|
@ -128,9 +140,13 @@ def train(
|
|||
if cfg.task == "dpo":
|
||||
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
||||
|
||||
trainer_wrapper = DPOTrainerWrapper(cfg, device=device, report_to=report_to)
|
||||
trainer_wrapper = DPOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
else:
|
||||
trainer_wrapper = SFTTrainerWrapper(cfg, device=device, report_to=report_to)
|
||||
trainer_wrapper = SFTTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper.setup(dataset)
|
||||
|
||||
# Train with live display and experiment tracking
|
||||
|
|
@ -173,6 +189,26 @@ def train(
|
|||
)
|
||||
|
||||
|
||||
def _resolve_deepspeed(deepspeed: str) -> str:
|
||||
"""Resolve DeepSpeed config: named preset or path to JSON file."""
|
||||
from soup_cli.utils.deepspeed import CONFIGS, write_deepspeed_config
|
||||
|
||||
# Named preset
|
||||
if deepspeed in CONFIGS:
|
||||
return write_deepspeed_config(deepspeed)
|
||||
|
||||
# Path to config file
|
||||
ds_path = Path(deepspeed)
|
||||
if ds_path.exists() and ds_path.suffix == ".json":
|
||||
return str(ds_path)
|
||||
|
||||
console.print(
|
||||
f"[red]Invalid DeepSpeed config: {deepspeed}[/]\n"
|
||||
f"Options: {', '.join(CONFIGS.keys())} or path to JSON file."
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
def _resolve_checkpoint(resume: str, output_dir: str, experiment_name: str = None) -> str:
|
||||
"""Resolve the checkpoint path from --resume argument.
|
||||
|
||||
|
|
|
|||
|
|
@ -21,10 +21,17 @@ class DPOTrainerWrapper:
|
|||
- rejected: the less preferred response
|
||||
"""
|
||||
|
||||
def __init__(self, config: SoupConfig, device: str = "cuda", report_to: str = "none"):
|
||||
def __init__(
|
||||
self,
|
||||
config: SoupConfig,
|
||||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.model = None
|
||||
self.ref_model = None
|
||||
self.tokenizer = None
|
||||
|
|
@ -154,6 +161,7 @@ class DPOTrainerWrapper:
|
|||
bf16=self.device == "cuda",
|
||||
report_to=self.report_to,
|
||||
remove_unused_columns=False,
|
||||
deepspeed=self.deepspeed_config,
|
||||
beta=tcfg.dpo_beta,
|
||||
max_length=cfg.data.max_length,
|
||||
max_prompt_length=cfg.data.max_length // 2,
|
||||
|
|
|
|||
|
|
@ -15,10 +15,17 @@ console = Console()
|
|||
class SFTTrainerWrapper:
|
||||
"""High-level wrapper that sets up model + tokenizer + trainer from SoupConfig."""
|
||||
|
||||
def __init__(self, config: SoupConfig, device: str = "cuda", report_to: str = "none"):
|
||||
def __init__(
|
||||
self,
|
||||
config: SoupConfig,
|
||||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -164,6 +171,7 @@ class SFTTrainerWrapper:
|
|||
bf16=self.device == "cuda",
|
||||
report_to=self.report_to,
|
||||
remove_unused_columns=False,
|
||||
deepspeed=self.deepspeed_config,
|
||||
)
|
||||
|
||||
# --- Trainer ---
|
||||
|
|
|
|||
|
|
@ -0,0 +1,117 @@
|
|||
"""DeepSpeed configuration templates for multi-GPU training."""
|
||||
|
||||
import copy
|
||||
import json
|
||||
import tempfile
|
||||
|
||||
# ZeRO Stage 2: splits optimizer states + gradients across GPUs
|
||||
ZERO_STAGE_2 = {
|
||||
"bf16": {"enabled": True},
|
||||
"zero_optimization": {
|
||||
"stage": 2,
|
||||
"offload_optimizer": {"device": "none"},
|
||||
"allgather_partitions": True,
|
||||
"allgather_bucket_size": 2e8,
|
||||
"overlap_comm": True,
|
||||
"reduce_scatter": True,
|
||||
"reduce_bucket_size": 2e8,
|
||||
"contiguous_gradients": True,
|
||||
},
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"gradient_clipping": "auto",
|
||||
"train_batch_size": "auto",
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"wall_clock_breakdown": False,
|
||||
}
|
||||
|
||||
# ZeRO Stage 3: splits model params + optimizer + gradients across GPUs
|
||||
ZERO_STAGE_3 = {
|
||||
"bf16": {"enabled": True},
|
||||
"zero_optimization": {
|
||||
"stage": 3,
|
||||
"offload_optimizer": {"device": "none"},
|
||||
"offload_param": {"device": "none"},
|
||||
"overlap_comm": True,
|
||||
"contiguous_gradients": True,
|
||||
"sub_group_size": 1e9,
|
||||
"reduce_bucket_size": "auto",
|
||||
"stage3_prefetch_bucket_size": "auto",
|
||||
"stage3_param_persistence_threshold": "auto",
|
||||
"stage3_max_live_parameters": 1e9,
|
||||
"stage3_max_reuse_distance": 1e9,
|
||||
"stage3_gather_16bit_weights_on_model_save": True,
|
||||
},
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"gradient_clipping": "auto",
|
||||
"train_batch_size": "auto",
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"wall_clock_breakdown": False,
|
||||
}
|
||||
|
||||
# ZeRO Stage 2 with CPU offload (for memory-constrained setups)
|
||||
ZERO_STAGE_2_OFFLOAD = {
|
||||
"bf16": {"enabled": True},
|
||||
"zero_optimization": {
|
||||
"stage": 2,
|
||||
"offload_optimizer": {"device": "cpu", "pin_memory": True},
|
||||
"allgather_partitions": True,
|
||||
"allgather_bucket_size": 2e8,
|
||||
"overlap_comm": True,
|
||||
"reduce_scatter": True,
|
||||
"reduce_bucket_size": 2e8,
|
||||
"contiguous_gradients": True,
|
||||
},
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"gradient_clipping": "auto",
|
||||
"train_batch_size": "auto",
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"wall_clock_breakdown": False,
|
||||
}
|
||||
|
||||
|
||||
CONFIGS = {
|
||||
"zero2": ZERO_STAGE_2,
|
||||
"zero3": ZERO_STAGE_3,
|
||||
"zero2_offload": ZERO_STAGE_2_OFFLOAD,
|
||||
}
|
||||
|
||||
|
||||
def get_deepspeed_config(stage: str = "zero2") -> dict:
|
||||
"""Get a DeepSpeed config dict by name."""
|
||||
if stage not in CONFIGS:
|
||||
raise ValueError(f"Unknown DeepSpeed config: {stage}. Options: {', '.join(CONFIGS.keys())}")
|
||||
return copy.deepcopy(CONFIGS[stage])
|
||||
|
||||
|
||||
def write_deepspeed_config(stage: str = "zero2") -> str:
|
||||
"""Write a DeepSpeed config to a temp file and return the path."""
|
||||
config = get_deepspeed_config(stage)
|
||||
tmp = tempfile.NamedTemporaryFile(
|
||||
mode="w", suffix=".json", prefix="ds_config_", delete=False
|
||||
)
|
||||
json.dump(config, tmp, indent=2)
|
||||
tmp.close()
|
||||
return tmp.name
|
||||
|
||||
|
||||
def detect_multi_gpu() -> dict:
|
||||
"""Detect multiple GPUs and return info."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
return {"gpu_count": 0, "gpus": []}
|
||||
|
||||
gpu_count = torch.cuda.device_count()
|
||||
gpus = []
|
||||
for idx in range(gpu_count):
|
||||
props = torch.cuda.get_device_properties(idx)
|
||||
gpus.append({
|
||||
"index": idx,
|
||||
"name": props.name,
|
||||
"memory_gb": props.total_memory / (1024 ** 3),
|
||||
})
|
||||
|
||||
return {"gpu_count": gpu_count, "gpus": gpus}
|
||||
except ImportError:
|
||||
return {"gpu_count": 0, "gpus": []}
|
||||
|
|
@ -0,0 +1,237 @@
|
|||
"""Tests for Multi-GPU / DeepSpeed support."""
|
||||
|
||||
import json
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class TestDeepSpeedConfigs:
|
||||
"""Test DeepSpeed configuration templates."""
|
||||
|
||||
def test_zero2_config_structure(self):
|
||||
"""ZeRO Stage 2 config should have correct structure."""
|
||||
from soup_cli.utils.deepspeed import get_deepspeed_config
|
||||
|
||||
config = get_deepspeed_config("zero2")
|
||||
assert config["zero_optimization"]["stage"] == 2
|
||||
assert config["bf16"]["enabled"] is True
|
||||
assert config["gradient_accumulation_steps"] == "auto"
|
||||
|
||||
def test_zero3_config_structure(self):
|
||||
"""ZeRO Stage 3 config should have correct structure."""
|
||||
from soup_cli.utils.deepspeed import get_deepspeed_config
|
||||
|
||||
config = get_deepspeed_config("zero3")
|
||||
assert config["zero_optimization"]["stage"] == 3
|
||||
assert config["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] is True
|
||||
|
||||
def test_zero2_offload_config(self):
|
||||
"""ZeRO Stage 2 with offload should enable CPU offloading."""
|
||||
from soup_cli.utils.deepspeed import get_deepspeed_config
|
||||
|
||||
config = get_deepspeed_config("zero2_offload")
|
||||
assert config["zero_optimization"]["stage"] == 2
|
||||
offload = config["zero_optimization"]["offload_optimizer"]
|
||||
assert offload["device"] == "cpu"
|
||||
assert offload["pin_memory"] is True
|
||||
|
||||
def test_invalid_config_name(self):
|
||||
"""Should raise ValueError for unknown config name."""
|
||||
from soup_cli.utils.deepspeed import get_deepspeed_config
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown DeepSpeed config"):
|
||||
get_deepspeed_config("zero99")
|
||||
|
||||
def test_get_config_returns_copy(self):
|
||||
"""Should return a copy, not the original."""
|
||||
from soup_cli.utils.deepspeed import get_deepspeed_config
|
||||
|
||||
config1 = get_deepspeed_config("zero2")
|
||||
config2 = get_deepspeed_config("zero2")
|
||||
config1["bf16"]["enabled"] = False
|
||||
assert config2["bf16"]["enabled"] is True
|
||||
|
||||
def test_all_configs_have_auto_fields(self):
|
||||
"""All configs should have 'auto' for batch sizes."""
|
||||
from soup_cli.utils.deepspeed import CONFIGS
|
||||
|
||||
for name, config in CONFIGS.items():
|
||||
assert config["train_batch_size"] == "auto", f"{name} missing auto train_batch_size"
|
||||
assert config["train_micro_batch_size_per_gpu"] == "auto", (
|
||||
f"{name} missing auto micro batch"
|
||||
)
|
||||
|
||||
|
||||
class TestWriteDeepSpeedConfig:
|
||||
"""Test writing DeepSpeed config to temp file."""
|
||||
|
||||
def test_write_creates_file(self):
|
||||
"""Should create a valid JSON file."""
|
||||
from soup_cli.utils.deepspeed import write_deepspeed_config
|
||||
|
||||
path = write_deepspeed_config("zero2")
|
||||
assert os.path.exists(path)
|
||||
|
||||
with open(path) as f:
|
||||
config = json.load(f)
|
||||
assert config["zero_optimization"]["stage"] == 2
|
||||
|
||||
# Cleanup
|
||||
os.unlink(path)
|
||||
|
||||
def test_write_file_is_valid_json(self):
|
||||
"""Written file should be parseable JSON."""
|
||||
from soup_cli.utils.deepspeed import write_deepspeed_config
|
||||
|
||||
for stage in ["zero2", "zero3", "zero2_offload"]:
|
||||
path = write_deepspeed_config(stage)
|
||||
with open(path) as f:
|
||||
config = json.load(f)
|
||||
assert "zero_optimization" in config
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
class TestDetectMultiGPU:
|
||||
"""Test multi-GPU detection."""
|
||||
|
||||
def test_detect_no_gpu(self):
|
||||
"""Should return 0 GPUs when CUDA not available."""
|
||||
from soup_cli.utils.deepspeed import detect_multi_gpu
|
||||
|
||||
with patch("torch.cuda.is_available", return_value=False):
|
||||
result = detect_multi_gpu()
|
||||
assert result["gpu_count"] == 0
|
||||
assert result["gpus"] == []
|
||||
|
||||
def test_detect_single_gpu(self):
|
||||
"""Should detect a single GPU."""
|
||||
from soup_cli.utils.deepspeed import detect_multi_gpu
|
||||
|
||||
mock_props = MagicMock()
|
||||
mock_props.name = "NVIDIA RTX 4090"
|
||||
mock_props.total_memory = 24 * (1024 ** 3) # 24GB
|
||||
|
||||
with patch("torch.cuda.is_available", return_value=True), \
|
||||
patch("torch.cuda.device_count", return_value=1), \
|
||||
patch("torch.cuda.get_device_properties", return_value=mock_props):
|
||||
result = detect_multi_gpu()
|
||||
assert result["gpu_count"] == 1
|
||||
assert len(result["gpus"]) == 1
|
||||
assert result["gpus"][0]["name"] == "NVIDIA RTX 4090"
|
||||
assert result["gpus"][0]["memory_gb"] == pytest.approx(24.0)
|
||||
|
||||
def test_detect_multiple_gpus(self):
|
||||
"""Should detect multiple GPUs."""
|
||||
from soup_cli.utils.deepspeed import detect_multi_gpu
|
||||
|
||||
mock_props = MagicMock()
|
||||
mock_props.name = "NVIDIA A100"
|
||||
mock_props.total_memory = 80 * (1024 ** 3)
|
||||
|
||||
with patch("torch.cuda.is_available", return_value=True), \
|
||||
patch("torch.cuda.device_count", return_value=4), \
|
||||
patch("torch.cuda.get_device_properties", return_value=mock_props):
|
||||
result = detect_multi_gpu()
|
||||
assert result["gpu_count"] == 4
|
||||
assert len(result["gpus"]) == 4
|
||||
|
||||
def test_detect_without_torch(self):
|
||||
"""Should handle missing torch gracefully."""
|
||||
from soup_cli.utils.deepspeed import detect_multi_gpu
|
||||
|
||||
with patch.dict("sys.modules", {"torch": None}):
|
||||
# Import error should be caught
|
||||
result = detect_multi_gpu()
|
||||
assert result["gpu_count"] == 0
|
||||
|
||||
|
||||
class TestResolveDeepSpeed:
|
||||
"""Test DeepSpeed config resolution in train command."""
|
||||
|
||||
def test_resolve_named_preset(self):
|
||||
"""Should resolve named presets like 'zero2'."""
|
||||
from soup_cli.commands.train import _resolve_deepspeed
|
||||
|
||||
path = _resolve_deepspeed("zero2")
|
||||
assert os.path.exists(path)
|
||||
|
||||
with open(path) as f:
|
||||
config = json.load(f)
|
||||
assert config["zero_optimization"]["stage"] == 2
|
||||
os.unlink(path)
|
||||
|
||||
def test_resolve_json_file(self, tmp_path):
|
||||
"""Should resolve path to JSON file."""
|
||||
from soup_cli.commands.train import _resolve_deepspeed
|
||||
|
||||
config_file = tmp_path / "ds_config.json"
|
||||
config_file.write_text(json.dumps({"zero_optimization": {"stage": 2}}))
|
||||
|
||||
result = _resolve_deepspeed(str(config_file))
|
||||
assert result == str(config_file)
|
||||
|
||||
def test_resolve_invalid_name(self):
|
||||
"""Should raise exit for invalid name."""
|
||||
from click.exceptions import Exit
|
||||
|
||||
from soup_cli.commands.train import _resolve_deepspeed
|
||||
|
||||
with pytest.raises(Exit):
|
||||
_resolve_deepspeed("invalid_config")
|
||||
|
||||
|
||||
class TestTrainerDeepSpeedParam:
|
||||
"""Test that trainers accept deepspeed_config parameter."""
|
||||
|
||||
def test_sft_trainer_accepts_deepspeed(self):
|
||||
"""SFTTrainerWrapper should accept deepspeed_config."""
|
||||
from soup_cli.config.schema import SoupConfig
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="test-model",
|
||||
data={"train": "test.jsonl"},
|
||||
)
|
||||
wrapper = SFTTrainerWrapper(cfg, device="cpu", deepspeed_config="/tmp/ds.json")
|
||||
assert wrapper.deepspeed_config == "/tmp/ds.json"
|
||||
|
||||
def test_dpo_trainer_accepts_deepspeed(self):
|
||||
"""DPOTrainerWrapper should accept deepspeed_config."""
|
||||
from soup_cli.config.schema import SoupConfig
|
||||
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="test-model",
|
||||
task="dpo",
|
||||
data={"train": "test.jsonl"},
|
||||
)
|
||||
wrapper = DPOTrainerWrapper(cfg, device="cpu", deepspeed_config="/tmp/ds.json")
|
||||
assert wrapper.deepspeed_config == "/tmp/ds.json"
|
||||
|
||||
def test_sft_trainer_default_no_deepspeed(self):
|
||||
"""SFTTrainerWrapper should default to no DeepSpeed."""
|
||||
from soup_cli.config.schema import SoupConfig
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="test-model",
|
||||
data={"train": "test.jsonl"},
|
||||
)
|
||||
wrapper = SFTTrainerWrapper(cfg, device="cpu")
|
||||
assert wrapper.deepspeed_config is None
|
||||
|
||||
|
||||
class TestTrainDeepSpeedFlag:
|
||||
"""Test --deepspeed flag in train command."""
|
||||
|
||||
def test_train_help_shows_deepspeed(self):
|
||||
"""Train help should mention --deepspeed option."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, ["train", "--help"])
|
||||
assert "--deepspeed" in result.output
|
||||
|
|
@ -0,0 +1,198 @@
|
|||
"""Tests for soup diff — model comparison command."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class TestCollectPrompts:
|
||||
"""Test prompt collection from files and arguments."""
|
||||
|
||||
def test_collect_from_text_file(self, tmp_path):
|
||||
"""Should read plain text prompts from a file."""
|
||||
from soup_cli.commands.diff import _collect_prompts
|
||||
|
||||
prompts_file = tmp_path / "prompts.txt"
|
||||
prompts_file.write_text("What is AI?\nExplain gravity.\nHello world.\n")
|
||||
|
||||
result = _collect_prompts(str(prompts_file), None)
|
||||
assert len(result) == 3
|
||||
assert result[0] == "What is AI?"
|
||||
assert result[1] == "Explain gravity."
|
||||
|
||||
def test_collect_from_jsonl_file(self, tmp_path):
|
||||
"""Should read prompts from JSONL with 'prompt' field."""
|
||||
from soup_cli.commands.diff import _collect_prompts
|
||||
|
||||
prompts_file = tmp_path / "prompts.jsonl"
|
||||
lines = [
|
||||
json.dumps({"prompt": "What is AI?"}),
|
||||
json.dumps({"prompt": "Explain gravity."}),
|
||||
]
|
||||
prompts_file.write_text("\n".join(lines))
|
||||
|
||||
result = _collect_prompts(str(prompts_file), None)
|
||||
assert len(result) == 2
|
||||
assert result[0] == "What is AI?"
|
||||
|
||||
def test_collect_from_args(self):
|
||||
"""Should collect prompts from CLI arguments."""
|
||||
from soup_cli.commands.diff import _collect_prompts
|
||||
|
||||
result = _collect_prompts(None, ["Hello!", "How are you?"])
|
||||
assert len(result) == 2
|
||||
assert result[0] == "Hello!"
|
||||
|
||||
def test_collect_combined(self, tmp_path):
|
||||
"""Should combine prompts from file and args."""
|
||||
from soup_cli.commands.diff import _collect_prompts
|
||||
|
||||
prompts_file = tmp_path / "prompts.txt"
|
||||
prompts_file.write_text("From file\n")
|
||||
|
||||
result = _collect_prompts(str(prompts_file), ["From args"])
|
||||
assert len(result) == 2
|
||||
|
||||
def test_collect_empty(self):
|
||||
"""Should return empty list if no prompts."""
|
||||
from soup_cli.commands.diff import _collect_prompts
|
||||
|
||||
result = _collect_prompts(None, None)
|
||||
assert result == []
|
||||
|
||||
def test_collect_skips_empty_lines(self, tmp_path):
|
||||
"""Should skip empty lines in prompt files."""
|
||||
from soup_cli.commands.diff import _collect_prompts
|
||||
|
||||
prompts_file = tmp_path / "prompts.txt"
|
||||
prompts_file.write_text("Line one\n\n\nLine two\n\n")
|
||||
|
||||
result = _collect_prompts(str(prompts_file), None)
|
||||
assert len(result) == 2
|
||||
|
||||
|
||||
class TestComputeMetrics:
|
||||
"""Test comparison metrics computation."""
|
||||
|
||||
def test_identical_responses(self):
|
||||
"""Identical responses should have 100% overlap."""
|
||||
from soup_cli.commands.diff import _compute_metrics
|
||||
|
||||
metrics = _compute_metrics("hello world", "hello world")
|
||||
assert metrics["word_overlap"] == pytest.approx(1.0)
|
||||
assert metrics["len_a"] == metrics["len_b"]
|
||||
|
||||
def test_completely_different(self):
|
||||
"""Completely different responses should have 0% overlap."""
|
||||
from soup_cli.commands.diff import _compute_metrics
|
||||
|
||||
metrics = _compute_metrics("hello world", "foo bar")
|
||||
assert metrics["word_overlap"] == pytest.approx(0.0)
|
||||
|
||||
def test_partial_overlap(self):
|
||||
"""Partial overlap should be between 0 and 1."""
|
||||
from soup_cli.commands.diff import _compute_metrics
|
||||
|
||||
metrics = _compute_metrics("hello world today", "hello world tomorrow")
|
||||
assert 0 < metrics["word_overlap"] < 1
|
||||
|
||||
def test_empty_responses(self):
|
||||
"""Empty responses should not crash."""
|
||||
from soup_cli.commands.diff import _compute_metrics
|
||||
|
||||
metrics = _compute_metrics("", "")
|
||||
assert metrics["len_a"] == 0
|
||||
assert metrics["len_b"] == 0
|
||||
assert metrics["word_overlap"] == pytest.approx(0.0)
|
||||
|
||||
def test_one_empty(self):
|
||||
"""One empty response should have 0% overlap."""
|
||||
from soup_cli.commands.diff import _compute_metrics
|
||||
|
||||
metrics = _compute_metrics("hello world", "")
|
||||
assert metrics["word_overlap"] == pytest.approx(0.0)
|
||||
|
||||
def test_word_counts(self):
|
||||
"""Should correctly count words."""
|
||||
from soup_cli.commands.diff import _compute_metrics
|
||||
|
||||
metrics = _compute_metrics("one two three", "a b")
|
||||
assert metrics["words_a"] == 3
|
||||
assert metrics["words_b"] == 2
|
||||
|
||||
|
||||
class TestDisplaySummary:
|
||||
"""Test summary display."""
|
||||
|
||||
def test_display_summary_no_crash(self):
|
||||
"""Display summary should not crash with valid data."""
|
||||
from soup_cli.commands.diff import _display_summary
|
||||
|
||||
results = [
|
||||
{
|
||||
"prompt": "test",
|
||||
"response_a": "hello",
|
||||
"response_b": "world",
|
||||
"metrics": {
|
||||
"len_a": 5, "len_b": 5,
|
||||
"words_a": 1, "words_b": 1,
|
||||
"word_overlap": 0.0,
|
||||
},
|
||||
}
|
||||
]
|
||||
# Should not raise
|
||||
_display_summary(results, "model_a", "model_b")
|
||||
|
||||
def test_display_summary_empty(self):
|
||||
"""Display summary should handle empty results."""
|
||||
from soup_cli.commands.diff import _display_summary
|
||||
|
||||
_display_summary([], "model_a", "model_b")
|
||||
|
||||
|
||||
class TestDiffCLI:
|
||||
"""Test diff CLI command."""
|
||||
|
||||
def test_model_a_not_found(self):
|
||||
"""Should fail if model A doesn't exist."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, [
|
||||
"diff",
|
||||
"--model-a", "/nonexistent/model_a",
|
||||
"--model-b", ".",
|
||||
"--prompt", "test",
|
||||
])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_model_b_not_found(self, tmp_path):
|
||||
"""Should fail if model B doesn't exist."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, [
|
||||
"diff",
|
||||
"--model-a", str(tmp_path),
|
||||
"--model-b", "/nonexistent/model_b",
|
||||
"--prompt", "test",
|
||||
])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_no_prompts_error(self, tmp_path):
|
||||
"""Should fail if no prompts provided."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, [
|
||||
"diff",
|
||||
"--model-a", str(tmp_path),
|
||||
"--model-b", str(tmp_path),
|
||||
])
|
||||
assert result.exit_code != 0
|
||||
|
|
@ -0,0 +1,224 @@
|
|||
"""Tests for soup data generate — synthetic data generation."""
|
||||
|
||||
import json
|
||||
|
||||
|
||||
class TestParseJsonArray:
|
||||
"""Test JSON array parsing from LLM output."""
|
||||
|
||||
def test_parse_clean_json_array(self):
|
||||
"""Should parse a clean JSON array."""
|
||||
from soup_cli.commands.generate import _parse_json_array
|
||||
|
||||
content = json.dumps([
|
||||
{"instruction": "What is AI?", "input": "", "output": "AI is..."},
|
||||
{"instruction": "Explain ML", "input": "", "output": "ML is..."},
|
||||
])
|
||||
result = _parse_json_array(content)
|
||||
assert len(result) == 2
|
||||
assert result[0]["instruction"] == "What is AI?"
|
||||
|
||||
def test_parse_json_with_markdown_fences(self):
|
||||
"""Should strip markdown code fences."""
|
||||
from soup_cli.commands.generate import _parse_json_array
|
||||
|
||||
content = '```json\n[{"instruction": "test", "input": "", "output": "ok"}]\n```'
|
||||
result = _parse_json_array(content)
|
||||
assert len(result) == 1
|
||||
assert result[0]["instruction"] == "test"
|
||||
|
||||
def test_parse_json_with_extra_text(self):
|
||||
"""Should extract JSON array from surrounding text."""
|
||||
from soup_cli.commands.generate import _parse_json_array
|
||||
|
||||
content = (
|
||||
'Here are the examples:\n'
|
||||
'[{"instruction": "a", "input": "", "output": "b"}]\nDone!'
|
||||
)
|
||||
result = _parse_json_array(content)
|
||||
assert len(result) == 1
|
||||
|
||||
def test_parse_empty_content(self):
|
||||
"""Should return empty list for empty content."""
|
||||
from soup_cli.commands.generate import _parse_json_array
|
||||
|
||||
result = _parse_json_array("")
|
||||
assert result == []
|
||||
|
||||
def test_parse_jsonl_fallback(self):
|
||||
"""Should fall back to line-by-line JSON parsing."""
|
||||
from soup_cli.commands.generate import _parse_json_array
|
||||
|
||||
content = (
|
||||
'{"instruction": "a", "input": "", "output": "b"}\n'
|
||||
'{"instruction": "c", "input": "", "output": "d"}'
|
||||
)
|
||||
result = _parse_json_array(content)
|
||||
assert len(result) == 2
|
||||
|
||||
def test_parse_invalid_json(self):
|
||||
"""Should return empty list for completely invalid JSON."""
|
||||
from soup_cli.commands.generate import _parse_json_array
|
||||
|
||||
result = _parse_json_array("this is not json at all")
|
||||
assert result == []
|
||||
|
||||
def test_parse_filters_non_dicts(self):
|
||||
"""Should filter out non-dict items."""
|
||||
from soup_cli.commands.generate import _parse_json_array
|
||||
|
||||
content = '[{"instruction": "a", "input": "", "output": "b"}, "string", 42]'
|
||||
result = _parse_json_array(content)
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
class TestValidateExample:
|
||||
"""Test example validation."""
|
||||
|
||||
def test_validate_alpaca_valid(self):
|
||||
"""Valid alpaca format should pass."""
|
||||
from soup_cli.commands.generate import _validate_example
|
||||
|
||||
example = {"instruction": "test", "input": "", "output": "ok"}
|
||||
assert _validate_example(example, "alpaca") is True
|
||||
|
||||
def test_validate_alpaca_missing_fields(self):
|
||||
"""Alpaca missing required fields should fail."""
|
||||
from soup_cli.commands.generate import _validate_example
|
||||
|
||||
assert _validate_example({"instruction": "test"}, "alpaca") is False
|
||||
assert _validate_example({"output": "test"}, "alpaca") is False
|
||||
|
||||
def test_validate_sharegpt_valid(self):
|
||||
"""Valid sharegpt format should pass."""
|
||||
from soup_cli.commands.generate import _validate_example
|
||||
|
||||
example = {
|
||||
"conversations": [
|
||||
{"from": "human", "value": "Hi"},
|
||||
{"from": "gpt", "value": "Hello!"},
|
||||
]
|
||||
}
|
||||
assert _validate_example(example, "sharegpt") is True
|
||||
|
||||
def test_validate_sharegpt_too_few(self):
|
||||
"""Sharegpt with fewer than 2 messages should fail."""
|
||||
from soup_cli.commands.generate import _validate_example
|
||||
|
||||
example = {"conversations": [{"from": "human", "value": "Hi"}]}
|
||||
assert _validate_example(example, "sharegpt") is False
|
||||
|
||||
def test_validate_chatml_valid(self):
|
||||
"""Valid chatml format should pass."""
|
||||
from soup_cli.commands.generate import _validate_example
|
||||
|
||||
example = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Hello!"},
|
||||
]
|
||||
}
|
||||
assert _validate_example(example, "chatml") is True
|
||||
|
||||
def test_validate_chatml_empty(self):
|
||||
"""Chatml with empty messages should fail."""
|
||||
from soup_cli.commands.generate import _validate_example
|
||||
|
||||
assert _validate_example({"messages": []}, "chatml") is False
|
||||
|
||||
def test_validate_unknown_format(self):
|
||||
"""Unknown format should fail."""
|
||||
from soup_cli.commands.generate import _validate_example
|
||||
|
||||
assert _validate_example({"data": "test"}, "unknown") is False
|
||||
|
||||
|
||||
class TestBuildGenerationPrompt:
|
||||
"""Test generation prompt building."""
|
||||
|
||||
def test_prompt_includes_format_spec(self):
|
||||
"""Prompt should include format specification."""
|
||||
from soup_cli.commands.generate import _build_generation_prompt
|
||||
|
||||
result = _build_generation_prompt("Create math questions", 5, "alpaca", [])
|
||||
assert "instruction" in result
|
||||
assert "output" in result
|
||||
assert "5" in result
|
||||
|
||||
def test_prompt_includes_seed_examples(self):
|
||||
"""Prompt should include seed examples when provided."""
|
||||
from soup_cli.commands.generate import _build_generation_prompt
|
||||
|
||||
seeds = [{"instruction": "example", "input": "", "output": "test"}]
|
||||
result = _build_generation_prompt("Create data", 3, "alpaca", seeds)
|
||||
assert "example" in result
|
||||
assert "seed" in result.lower()
|
||||
|
||||
def test_prompt_sharegpt_format(self):
|
||||
"""Prompt should describe sharegpt format correctly."""
|
||||
from soup_cli.commands.generate import _build_generation_prompt
|
||||
|
||||
result = _build_generation_prompt("Create chats", 3, "sharegpt", [])
|
||||
assert "conversations" in result
|
||||
|
||||
def test_prompt_chatml_format(self):
|
||||
"""Prompt should describe chatml format correctly."""
|
||||
from soup_cli.commands.generate import _build_generation_prompt
|
||||
|
||||
result = _build_generation_prompt("Create chats", 3, "chatml", [])
|
||||
assert "messages" in result
|
||||
|
||||
|
||||
class TestRowToText:
|
||||
"""Test row to text conversion for dedup."""
|
||||
|
||||
def test_row_to_text_basic(self):
|
||||
"""Should concatenate all values."""
|
||||
from soup_cli.commands.generate import _row_to_text
|
||||
|
||||
row = {"instruction": "What is AI?", "output": "AI is..."}
|
||||
text = _row_to_text(row)
|
||||
assert "What is AI?" in text
|
||||
assert "AI is..." in text
|
||||
|
||||
def test_row_to_text_empty_values(self):
|
||||
"""Should skip empty values."""
|
||||
from soup_cli.commands.generate import _row_to_text
|
||||
|
||||
row = {"instruction": "test", "input": "", "output": "ok"}
|
||||
text = _row_to_text(row)
|
||||
assert "test" in text
|
||||
|
||||
|
||||
class TestGenerateCLI:
|
||||
"""Test CLI integration for generate command."""
|
||||
|
||||
def test_invalid_format_rejected(self):
|
||||
"""Should reject invalid format."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, [
|
||||
"data", "generate",
|
||||
"--prompt", "test",
|
||||
"--format", "invalid_format",
|
||||
"--count", "1",
|
||||
])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_invalid_provider_rejected(self):
|
||||
"""Should reject invalid provider."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, [
|
||||
"data", "generate",
|
||||
"--prompt", "test",
|
||||
"--provider", "invalid_provider",
|
||||
"--count", "1",
|
||||
])
|
||||
assert result.exit_code != 0
|
||||
|
|
@ -0,0 +1,257 @@
|
|||
"""Tests for soup serve — inference server command."""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class TestServeValidation:
|
||||
"""Test serve command argument validation."""
|
||||
|
||||
def test_model_path_not_found(self):
|
||||
"""serve should fail if model path doesn't exist."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, ["serve", "--model", "/nonexistent/path"])
|
||||
assert result.exit_code != 0
|
||||
assert "not found" in result.output.lower() or "not installed" in result.output.lower()
|
||||
|
||||
def test_fastapi_import_error(self):
|
||||
"""serve should fail gracefully if FastAPI not installed."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
with patch.dict("sys.modules", {"fastapi": None}):
|
||||
result = runner.invoke(app, ["serve", "--model", "."])
|
||||
# Either import error or model not found
|
||||
assert result.exit_code != 0
|
||||
|
||||
|
||||
class TestCreateApp:
|
||||
"""Test the FastAPI app creation."""
|
||||
|
||||
def test_create_app_returns_fastapi_instance(self):
|
||||
"""_create_app should return a FastAPI app with correct endpoints."""
|
||||
try:
|
||||
from fastapi import FastAPI
|
||||
except ImportError:
|
||||
pytest.skip("FastAPI not installed")
|
||||
|
||||
from soup_cli.commands.serve import _create_app
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_tokenizer = MagicMock()
|
||||
|
||||
app = _create_app(
|
||||
model_obj=mock_model,
|
||||
tokenizer=mock_tokenizer,
|
||||
device="cpu",
|
||||
model_name="test-model",
|
||||
max_tokens_default=256,
|
||||
)
|
||||
|
||||
assert isinstance(app, FastAPI)
|
||||
|
||||
# Check routes exist
|
||||
routes = [route.path for route in app.routes]
|
||||
assert "/health" in routes
|
||||
assert "/v1/models" in routes
|
||||
assert "/v1/chat/completions" in routes
|
||||
|
||||
def test_health_endpoint(self):
|
||||
"""Health endpoint should return ok status."""
|
||||
try:
|
||||
from fastapi.testclient import TestClient
|
||||
except ImportError:
|
||||
pytest.skip("FastAPI not installed")
|
||||
|
||||
from soup_cli.commands.serve import _create_app
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_tokenizer = MagicMock()
|
||||
|
||||
app = _create_app(
|
||||
model_obj=mock_model,
|
||||
tokenizer=mock_tokenizer,
|
||||
device="cpu",
|
||||
model_name="test-model",
|
||||
max_tokens_default=256,
|
||||
)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/health")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "ok"
|
||||
assert data["model"] == "test-model"
|
||||
assert data["device"] == "cpu"
|
||||
|
||||
def test_models_endpoint(self):
|
||||
"""Models endpoint should list the loaded model."""
|
||||
try:
|
||||
from fastapi.testclient import TestClient
|
||||
except ImportError:
|
||||
pytest.skip("FastAPI not installed")
|
||||
|
||||
from soup_cli.commands.serve import _create_app
|
||||
|
||||
app = _create_app(
|
||||
model_obj=MagicMock(),
|
||||
tokenizer=MagicMock(),
|
||||
device="cpu",
|
||||
model_name="my-model",
|
||||
max_tokens_default=256,
|
||||
)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/v1/models")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["object"] == "list"
|
||||
assert len(data["data"]) == 1
|
||||
assert data["data"][0]["id"] == "my-model"
|
||||
|
||||
def test_chat_completions_endpoint(self):
|
||||
"""Chat completions should return OpenAI-compatible response."""
|
||||
try:
|
||||
from fastapi.testclient import TestClient
|
||||
except ImportError:
|
||||
pytest.skip("FastAPI not installed")
|
||||
|
||||
from soup_cli.commands.serve import _create_app
|
||||
|
||||
with patch("soup_cli.commands.serve._generate_response") as mock_gen:
|
||||
mock_gen.return_value = ("Hello there!", 10, 5)
|
||||
|
||||
app = _create_app(
|
||||
model_obj=MagicMock(),
|
||||
tokenizer=MagicMock(),
|
||||
device="cpu",
|
||||
model_name="test-model",
|
||||
max_tokens_default=256,
|
||||
)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"model": "test-model",
|
||||
"messages": [{"role": "user", "content": "Hi"}],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["object"] == "chat.completion"
|
||||
assert data["choices"][0]["message"]["role"] == "assistant"
|
||||
assert data["choices"][0]["message"]["content"] == "Hello there!"
|
||||
assert data["choices"][0]["finish_reason"] == "stop"
|
||||
assert data["usage"]["prompt_tokens"] == 10
|
||||
assert data["usage"]["completion_tokens"] == 5
|
||||
|
||||
def test_chat_completions_streaming(self):
|
||||
"""Streaming should return SSE events."""
|
||||
try:
|
||||
from fastapi.testclient import TestClient
|
||||
except ImportError:
|
||||
pytest.skip("FastAPI not installed")
|
||||
|
||||
from soup_cli.commands.serve import _create_app
|
||||
|
||||
with patch("soup_cli.commands.serve._generate_response") as mock_gen:
|
||||
mock_gen.return_value = ("Hello world", 10, 5)
|
||||
|
||||
app = _create_app(
|
||||
model_obj=MagicMock(),
|
||||
tokenizer=MagicMock(),
|
||||
device="cpu",
|
||||
model_name="test-model",
|
||||
max_tokens_default=256,
|
||||
)
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"model": "test-model",
|
||||
"messages": [{"role": "user", "content": "Hi"}],
|
||||
"stream": True,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
text = response.text
|
||||
assert "data:" in text
|
||||
assert "[DONE]" in text
|
||||
|
||||
|
||||
class TestStreamResponse:
|
||||
"""Test SSE streaming."""
|
||||
|
||||
def test_stream_response_yields_chunks(self):
|
||||
"""_stream_response should yield SSE events."""
|
||||
from soup_cli.commands.serve import _stream_response
|
||||
|
||||
with patch("soup_cli.commands.serve._generate_response") as mock_gen:
|
||||
mock_gen.return_value = ("Hello world", 5, 2)
|
||||
|
||||
chunks = list(_stream_response(
|
||||
model=MagicMock(),
|
||||
tokenizer=MagicMock(),
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
max_tokens=256,
|
||||
temperature=0.7,
|
||||
top_p=0.9,
|
||||
model_name="test",
|
||||
))
|
||||
|
||||
# Should have word chunks + final chunk + [DONE]
|
||||
assert len(chunks) >= 3
|
||||
assert chunks[-1] == "data: [DONE]\n\n"
|
||||
|
||||
# First chunk should be valid SSE JSON
|
||||
first_data = json.loads(chunks[0].replace("data: ", "").strip())
|
||||
assert first_data["object"] == "chat.completion.chunk"
|
||||
assert first_data["choices"][0]["delta"]["content"] == "Hello"
|
||||
|
||||
|
||||
class TestDetectBaseModel:
|
||||
"""Test base model detection from adapter config."""
|
||||
|
||||
def test_detect_base_model_from_adapter(self, tmp_path):
|
||||
"""Should read base model from adapter_config.json."""
|
||||
from soup_cli.commands.serve import _detect_base_model
|
||||
|
||||
config_file = tmp_path / "adapter_config.json"
|
||||
config_file.write_text(json.dumps({
|
||||
"base_model_name_or_path": "meta-llama/Llama-3.1-8B"
|
||||
}))
|
||||
|
||||
result = _detect_base_model(config_file)
|
||||
assert result == "meta-llama/Llama-3.1-8B"
|
||||
|
||||
def test_detect_base_model_missing_key(self, tmp_path):
|
||||
"""Should return None if key is missing."""
|
||||
from soup_cli.commands.serve import _detect_base_model
|
||||
|
||||
config_file = tmp_path / "adapter_config.json"
|
||||
config_file.write_text(json.dumps({"other_key": "value"}))
|
||||
|
||||
result = _detect_base_model(config_file)
|
||||
assert result is None
|
||||
|
||||
def test_detect_base_model_invalid_json(self, tmp_path):
|
||||
"""Should return None for invalid JSON."""
|
||||
from soup_cli.commands.serve import _detect_base_model
|
||||
|
||||
config_file = tmp_path / "adapter_config.json"
|
||||
config_file.write_text("not json")
|
||||
|
||||
result = _detect_base_model(config_file)
|
||||
assert result is None
|
||||
|
|
@ -0,0 +1,265 @@
|
|||
"""Tests for soup sweep — hyperparameter search."""
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class TestParseSweepParams:
|
||||
"""Test sweep parameter parsing."""
|
||||
|
||||
def test_parse_simple_float_params(self):
|
||||
"""Should parse float values like learning rates."""
|
||||
from soup_cli.commands.sweep import _parse_sweep_params
|
||||
|
||||
result = _parse_sweep_params(["lr=1e-5,2e-5,5e-5"])
|
||||
assert "lr" in result
|
||||
assert len(result["lr"]) == 3
|
||||
assert result["lr"][0] == pytest.approx(1e-5)
|
||||
assert result["lr"][1] == pytest.approx(2e-5)
|
||||
assert result["lr"][2] == pytest.approx(5e-5)
|
||||
|
||||
def test_parse_int_params(self):
|
||||
"""Should parse integer values."""
|
||||
from soup_cli.commands.sweep import _parse_sweep_params
|
||||
|
||||
result = _parse_sweep_params(["lora_r=8,16,32"])
|
||||
assert result["lora_r"] == [8, 16, 32]
|
||||
|
||||
def test_parse_string_params(self):
|
||||
"""Should parse string values."""
|
||||
from soup_cli.commands.sweep import _parse_sweep_params
|
||||
|
||||
result = _parse_sweep_params(["optimizer=adamw_torch,sgd"])
|
||||
assert result["optimizer"] == ["adamw_torch", "sgd"]
|
||||
|
||||
def test_parse_bool_params(self):
|
||||
"""Should parse boolean values."""
|
||||
from soup_cli.commands.sweep import _parse_sweep_params
|
||||
|
||||
result = _parse_sweep_params(["some_flag=true,false"])
|
||||
assert result["some_flag"] == [True, False]
|
||||
|
||||
def test_parse_multiple_params(self):
|
||||
"""Should parse multiple parameter strings."""
|
||||
from soup_cli.commands.sweep import _parse_sweep_params
|
||||
|
||||
result = _parse_sweep_params(["lr=1e-5,2e-5", "epochs=2,3"])
|
||||
assert "lr" in result
|
||||
assert "epochs" in result
|
||||
assert len(result["lr"]) == 2
|
||||
assert len(result["epochs"]) == 2
|
||||
|
||||
def test_parse_invalid_param_no_equals(self):
|
||||
"""Should skip params without equals sign."""
|
||||
from soup_cli.commands.sweep import _parse_sweep_params
|
||||
|
||||
result = _parse_sweep_params(["invalid_param"])
|
||||
assert len(result) == 0
|
||||
|
||||
def test_parse_none_value(self):
|
||||
"""Should parse 'none' as None."""
|
||||
from soup_cli.commands.sweep import _parse_sweep_params
|
||||
|
||||
result = _parse_sweep_params(["target_modules=auto,none"])
|
||||
assert result["target_modules"] == ["auto", None]
|
||||
|
||||
|
||||
class TestParseValue:
|
||||
"""Test individual value parsing."""
|
||||
|
||||
def test_parse_int(self):
|
||||
from soup_cli.commands.sweep import _parse_value
|
||||
|
||||
assert _parse_value("42") == 42
|
||||
|
||||
def test_parse_float(self):
|
||||
from soup_cli.commands.sweep import _parse_value
|
||||
|
||||
assert _parse_value("3.14") == pytest.approx(3.14)
|
||||
|
||||
def test_parse_scientific_notation(self):
|
||||
from soup_cli.commands.sweep import _parse_value
|
||||
|
||||
assert _parse_value("1e-5") == pytest.approx(1e-5)
|
||||
|
||||
def test_parse_bool_true(self):
|
||||
from soup_cli.commands.sweep import _parse_value
|
||||
|
||||
assert _parse_value("true") is True
|
||||
|
||||
def test_parse_bool_false(self):
|
||||
from soup_cli.commands.sweep import _parse_value
|
||||
|
||||
assert _parse_value("false") is False
|
||||
|
||||
def test_parse_none(self):
|
||||
from soup_cli.commands.sweep import _parse_value
|
||||
|
||||
assert _parse_value("none") is None
|
||||
|
||||
def test_parse_string(self):
|
||||
from soup_cli.commands.sweep import _parse_value
|
||||
|
||||
assert _parse_value("adamw_torch") == "adamw_torch"
|
||||
|
||||
|
||||
class TestGenerateCombinations:
|
||||
"""Test parameter combination generation."""
|
||||
|
||||
def test_grid_search(self):
|
||||
"""Grid search should generate all combinations."""
|
||||
from soup_cli.commands.sweep import _generate_combinations
|
||||
|
||||
params = {"lr": [1e-5, 2e-5], "epochs": [2, 3]}
|
||||
combos = _generate_combinations(params, "grid", None)
|
||||
assert len(combos) == 4 # 2 x 2
|
||||
|
||||
def test_grid_search_single_param(self):
|
||||
"""Grid with one param should equal param count."""
|
||||
from soup_cli.commands.sweep import _generate_combinations
|
||||
|
||||
params = {"lr": [1e-5, 2e-5, 5e-5]}
|
||||
combos = _generate_combinations(params, "grid", None)
|
||||
assert len(combos) == 3
|
||||
|
||||
def test_grid_search_max_runs(self):
|
||||
"""Grid with max_runs should truncate."""
|
||||
from soup_cli.commands.sweep import _generate_combinations
|
||||
|
||||
params = {"lr": [1e-5, 2e-5], "epochs": [2, 3]}
|
||||
combos = _generate_combinations(params, "grid", max_runs=2)
|
||||
assert len(combos) == 2
|
||||
|
||||
def test_random_search(self):
|
||||
"""Random search should respect max_runs."""
|
||||
from soup_cli.commands.sweep import _generate_combinations
|
||||
|
||||
params = {"lr": [1e-5, 2e-5, 5e-5], "epochs": [2, 3, 5]}
|
||||
combos = _generate_combinations(params, "random", max_runs=3)
|
||||
assert len(combos) == 3
|
||||
|
||||
def test_random_search_no_duplicates(self):
|
||||
"""Random search should not produce duplicates."""
|
||||
from soup_cli.commands.sweep import _generate_combinations
|
||||
|
||||
params = {"lr": [1e-5, 2e-5], "epochs": [2, 3]}
|
||||
combos = _generate_combinations(params, "random", max_runs=4)
|
||||
combo_tuples = [tuple(sorted(c.items())) for c in combos]
|
||||
assert len(combo_tuples) == len(set(combo_tuples))
|
||||
|
||||
def test_random_search_defaults_to_10(self):
|
||||
"""Random search without max_runs should default to min(total, 10)."""
|
||||
from soup_cli.commands.sweep import _generate_combinations
|
||||
|
||||
params = {"lr": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]}
|
||||
combos = _generate_combinations(params, "random", max_runs=None)
|
||||
assert len(combos) == 10
|
||||
|
||||
|
||||
class TestSetNestedParam:
|
||||
"""Test nested parameter setting with shortcuts."""
|
||||
|
||||
def test_set_shortcut_lr(self):
|
||||
"""'lr' shortcut should set training.lr."""
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {"training": {"lr": 1e-5, "epochs": 3}}
|
||||
_set_nested_param(config, "lr", 2e-5)
|
||||
assert config["training"]["lr"] == pytest.approx(2e-5)
|
||||
|
||||
def test_set_shortcut_epochs(self):
|
||||
"""'epochs' shortcut should set training.epochs."""
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {"training": {"epochs": 3}}
|
||||
_set_nested_param(config, "epochs", 5)
|
||||
assert config["training"]["epochs"] == 5
|
||||
|
||||
def test_set_shortcut_lora_r(self):
|
||||
"""'lora_r' shortcut should set training.lora.r."""
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {"training": {"lora": {"r": 64}}}
|
||||
_set_nested_param(config, "lora_r", 32)
|
||||
assert config["training"]["lora"]["r"] == 32
|
||||
|
||||
def test_set_dot_notation(self):
|
||||
"""Dot notation should work for custom paths."""
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {"training": {"weight_decay": 0.01}}
|
||||
_set_nested_param(config, "training.weight_decay", 0.1)
|
||||
assert config["training"]["weight_decay"] == pytest.approx(0.1)
|
||||
|
||||
def test_set_creates_missing_keys(self):
|
||||
"""Should create intermediate keys if missing."""
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {}
|
||||
_set_nested_param(config, "training.lr", 1e-5)
|
||||
assert config["training"]["lr"] == pytest.approx(1e-5)
|
||||
|
||||
|
||||
class TestSweepCLI:
|
||||
"""Test sweep CLI command."""
|
||||
|
||||
def test_sweep_dry_run(self, tmp_path):
|
||||
"""Dry run should show plan without executing."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
# Create a minimal config
|
||||
config_file = tmp_path / "soup.yaml"
|
||||
config_file.write_text(
|
||||
"base: test-model\n"
|
||||
"data:\n"
|
||||
" train: ./data.jsonl\n"
|
||||
)
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, [
|
||||
"sweep",
|
||||
"--config", str(config_file),
|
||||
"--param", "lr=1e-5,2e-5",
|
||||
"--dry-run",
|
||||
])
|
||||
assert result.exit_code == 0
|
||||
assert "Sweep Plan" in result.output
|
||||
assert "lr" in result.output
|
||||
|
||||
def test_sweep_config_not_found(self):
|
||||
"""Should fail if config doesn't exist."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, [
|
||||
"sweep",
|
||||
"--config", "/nonexistent/soup.yaml",
|
||||
"--param", "lr=1e-5,2e-5",
|
||||
])
|
||||
assert result.exit_code != 0
|
||||
|
||||
def test_sweep_invalid_strategy(self, tmp_path):
|
||||
"""Should fail for invalid strategy."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
config_file = tmp_path / "soup.yaml"
|
||||
config_file.write_text(
|
||||
"base: test-model\n"
|
||||
"data:\n"
|
||||
" train: ./data.jsonl\n"
|
||||
)
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, [
|
||||
"sweep",
|
||||
"--config", str(config_file),
|
||||
"--param", "lr=1e-5",
|
||||
"--strategy", "bayesian",
|
||||
])
|
||||
assert result.exit_code != 0
|
||||
Loading…
Reference in New Issue