mirror of https://github.com/razor-ai/soup.git
Add multimodal vision fine-tuning support (Phase 6) — v0.5.0
- Add `modality: vision` config option for vision-language model training - Add LLaVA and ShareGPT4V data format detection and conversion - Add `image_dir` field in DataConfig for resolving image paths - Add vision model loading via AutoModelForVision2Seq + AutoProcessor in SFT trainer - Add `soup init --template vision` with LLaMA-3.2-Vision config - Add image statistics display in `soup data inspect` for vision datasets - Add Pillow as optional `vision` extra dependency - Add Pillow to `soup doctor` dependency checks - 51 new tests (455 total), ruff clean Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alpamys
parent
3b46e96797
commit
5b7ad3c358
|
|
@ -39,11 +39,11 @@ soup train --config soup.yaml
|
|||
→ save LoRA adapter to output/
|
||||
```
|
||||
|
||||
**Config system:** `config/schema.py` is the single source of truth. All YAML fields are validated by Pydantic models (`SoupConfig` → `TrainingConfig` → `LoraConfig`, `DataConfig`). Templates (chat/code/medical) live as YAML strings in this file.
|
||||
**Config system:** `config/schema.py` is the single source of truth. All YAML fields are validated by Pydantic models (`SoupConfig` → `TrainingConfig` → `LoraConfig`, `DataConfig`). Templates (chat/code/medical/reasoning/vision) live as YAML strings in this file. `SoupConfig.modality` supports `text` (default) and `vision` for multimodal fine-tuning. `DataConfig.format` includes `llava` and `sharegpt4v` for vision datasets, and `DataConfig.image_dir` specifies the base directory for resolving image paths.
|
||||
|
||||
**Data pipeline:** `data/loader.py` handles local files (JSONL/JSON/CSV/Parquet) and HuggingFace datasets. `data/formats.py` auto-detects and normalizes alpaca/sharegpt/chatml formats into a unified `{"messages": [...]}` structure. Also supports reverse conversion via `messages_to_format()`.
|
||||
**Data pipeline:** `data/loader.py` handles local files (JSONL/JSON/CSV/Parquet) and HuggingFace datasets. `data/formats.py` auto-detects and normalizes alpaca/sharegpt/chatml/llava/sharegpt4v formats into a unified `{"messages": [...]}` structure. Vision formats (llava, sharegpt4v) also include an `"image"` key. Also supports reverse conversion via `messages_to_format()`. Vision datasets get image path validation via `_validate_vision_images()`.
|
||||
|
||||
**Trainer:** `trainer/sft.py` (`SFTTrainerWrapper`), `trainer/dpo.py` (`DPOTrainerWrapper`), and `trainer/grpo.py` (`GRPOTrainerWrapper`) wrap HuggingFace's SFTTrainer/DPOTrainer/GRPOTrainer with auto quantization (BitsAndBytes), LoRA (PEFT), and batch size estimation. Heavy ML imports are lazy (inside methods) so CLI stays fast for non-training commands. All trainers enable Rich progress bars for HuggingFace Hub model downloads via `_enable_hf_transfer_progress()`. All trainers support `backend: unsloth` — when enabled, model loading is delegated to `utils/unsloth.py` which uses `unsloth.FastLanguageModel` for 2-5x faster training.
|
||||
**Trainer:** `trainer/sft.py` (`SFTTrainerWrapper`), `trainer/dpo.py` (`DPOTrainerWrapper`), and `trainer/grpo.py` (`GRPOTrainerWrapper`) wrap HuggingFace's SFTTrainer/DPOTrainer/GRPOTrainer with auto quantization (BitsAndBytes), LoRA (PEFT), and batch size estimation. Heavy ML imports are lazy (inside methods) so CLI stays fast for non-training commands. All trainers enable Rich progress bars for HuggingFace Hub model downloads via `_enable_hf_transfer_progress()`. All trainers support `backend: unsloth` — when enabled, model loading is delegated to `utils/unsloth.py` which uses `unsloth.FastLanguageModel` for 2-5x faster training. SFT trainer supports `modality: vision` — uses `AutoModelForVision2Seq` + `AutoProcessor` for vision-language models (LLaMA-Vision, Qwen2-VL, Pixtral).
|
||||
|
||||
**Unsloth backend:** `utils/unsloth.py` provides `is_unsloth_available()`, `get_unsloth_version()`, and `load_model_and_tokenizer()`. The latter uses `FastLanguageModel.from_pretrained()` + `FastLanguageModel.get_peft_model()` which handles quantization, LoRA patching, and kernel optimization internally. Config: `backend: unsloth` (default: `transformers`). Install: `pip install 'soup-cli[fast]'`. Auto-detection: if unsloth is installed but not enabled, `commands/train.py` shows a hint.
|
||||
|
||||
|
|
@ -167,3 +167,4 @@ Test suite lives in `tests/`:
|
|||
| `test_grpo.py` | GRPO config, rewards, data prep, template, sweep shortcuts |
|
||||
| `test_progress.py` | Rich download progress bar, `_enable_hf_transfer_progress` |
|
||||
| `test_unsloth.py` | Unsloth backend config, detection, trainer integration, templates |
|
||||
| `test_vision.py` | Vision modality config, LLaVA/ShareGPT4V formats, loader, trainer, templates |
|
||||
|
|
|
|||
63
README.md
63
README.md
|
|
@ -68,6 +68,7 @@ soup init --template chat # conversational fine-tune
|
|||
soup init --template code # code generation
|
||||
soup init --template medical # domain expert
|
||||
soup init --template reasoning # GRPO reasoning training
|
||||
soup init --template vision # vision/multimodal fine-tune
|
||||
```
|
||||
|
||||
### 3. Train
|
||||
|
|
@ -157,6 +158,55 @@ Works with all training tasks: SFT, DPO, and GRPO. If unsloth is installed but n
|
|||
|
||||
> **Tip:** Soup auto-detects unsloth. When installed, you'll see a hint during `soup train` if you haven't enabled it yet.
|
||||
|
||||
## Vision / Multimodal Fine-tuning
|
||||
|
||||
Fine-tune vision-language models (LLaMA-3.2-Vision, Qwen2-VL, Pixtral) on image+text data:
|
||||
|
||||
```bash
|
||||
# Install vision support
|
||||
pip install 'soup-cli[vision]'
|
||||
|
||||
# Create a vision config
|
||||
soup init --template vision
|
||||
|
||||
# Train
|
||||
soup train --config soup.yaml
|
||||
```
|
||||
|
||||
```yaml
|
||||
base: meta-llama/Llama-3.2-11B-Vision-Instruct
|
||||
task: sft
|
||||
modality: vision
|
||||
|
||||
data:
|
||||
train: ./data/vision_train.jsonl
|
||||
format: llava
|
||||
image_dir: ./data/images
|
||||
val_split: 0.1
|
||||
|
||||
training:
|
||||
epochs: 3
|
||||
lr: 1e-5
|
||||
quantization: 4bit
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
```
|
||||
|
||||
**Supported vision data formats:**
|
||||
|
||||
**LLaVA:**
|
||||
```json
|
||||
{"image": "photo.jpg", "conversations": [{"from": "human", "value": "<image>\nDescribe this image."}, {"from": "gpt", "value": "A cat on a mat."}]}
|
||||
```
|
||||
|
||||
**ShareGPT4V:**
|
||||
```json
|
||||
{"image": "chart.png", "conversations": [{"from": "human", "value": "<image>\nWhat does this show?"}, {"from": "gpt", "value": "Quarterly revenue."}]}
|
||||
```
|
||||
|
||||
`soup data inspect` automatically shows image statistics (count, formats, missing files) for vision datasets.
|
||||
|
||||
## DPO Training
|
||||
|
||||
Train with preference data using Direct Preference Optimization:
|
||||
|
|
@ -486,6 +536,16 @@ Soup supports these formats (auto-detected). Files can be JSONL, JSON, CSV, or P
|
|||
{"prompt": "Explain gravity", "chosen": "Gravity is a force...", "rejected": "I don't know"}
|
||||
```
|
||||
|
||||
**LLaVA (vision):**
|
||||
```json
|
||||
{"image": "photo.jpg", "conversations": [{"from": "human", "value": "<image>\nDescribe this."}, {"from": "gpt", "value": "A cat."}]}
|
||||
```
|
||||
|
||||
**ShareGPT4V (vision):**
|
||||
```json
|
||||
{"image": "chart.png", "conversations": [{"from": "human", "value": "<image>\nExplain this chart."}, {"from": "gpt", "value": "Revenue growth."}]}
|
||||
```
|
||||
|
||||
## Data Tools
|
||||
|
||||
```bash
|
||||
|
|
@ -544,7 +604,7 @@ soup eval --model ./output --benchmarks mmlu --run-id run_20260223_143052_a1b2
|
|||
## All Commands
|
||||
|
||||
```
|
||||
soup init [--template chat|code|medical|reasoning] Create config
|
||||
soup init [--template chat|code|medical|reasoning|vision] Create config
|
||||
soup train --config soup.yaml Start training
|
||||
soup chat --model ./output Interactive chat
|
||||
soup push --model ./output --repo user/name Upload to HuggingFace
|
||||
|
|
@ -580,6 +640,7 @@ soup --verbose <command> Full traceback on errors
|
|||
|
||||
| Extra | Install | What it adds |
|
||||
|---|---|---|
|
||||
| `vision` | `pip install 'soup-cli[vision]'` | Vision/multimodal fine-tuning (Pillow) |
|
||||
| `fast` | `pip install 'soup-cli[fast]'` | Unsloth backend (2-5x faster, -80% VRAM) |
|
||||
| `serve` | `pip install 'soup-cli[serve]'` | Inference server (FastAPI + uvicorn) |
|
||||
| `data` | `pip install 'soup-cli[data]'` | Deduplication (MinHash via datasketch) |
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.4.3"
|
||||
version = "0.5.0"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
|
@ -47,6 +47,7 @@ serve = ["fastapi>=0.104.0", "uvicorn>=0.24.0"]
|
|||
generate = ["httpx>=0.24.0"]
|
||||
deepspeed = ["deepspeed>=0.12.0"]
|
||||
fast = ["unsloth>=2024.8"]
|
||||
vision = ["Pillow>=9.0.0"]
|
||||
|
||||
[project.scripts]
|
||||
soup = "soup_cli.cli:run"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.4.3"
|
||||
__version__ = "0.5.0"
|
||||
|
|
|
|||
|
|
@ -46,6 +46,9 @@ def inspect(
|
|||
stats_table.add_row("Duplicates", str(result["duplicates"]))
|
||||
console.print(stats_table)
|
||||
|
||||
# Vision stats (if dataset contains images)
|
||||
_show_vision_stats(data)
|
||||
|
||||
# Print sample rows
|
||||
if rows > 0 and len(data) > 0:
|
||||
console.print(f"\n[bold]Sample rows ({min(rows, len(data))}):[/]")
|
||||
|
|
@ -364,6 +367,44 @@ def stats(
|
|||
)
|
||||
|
||||
|
||||
def _show_vision_stats(data: list[dict]) -> None:
|
||||
"""Show image statistics if dataset contains image fields."""
|
||||
if not data:
|
||||
return
|
||||
|
||||
# Check if this is a vision dataset
|
||||
sample = data[0]
|
||||
if "image" not in sample:
|
||||
return
|
||||
|
||||
total = len(data)
|
||||
has_image = sum(1 for row in data if row.get("image"))
|
||||
missing_image = total - has_image
|
||||
|
||||
# Collect image file info
|
||||
extensions: dict[str, int] = {}
|
||||
existing = 0
|
||||
for row in data:
|
||||
img_path = row.get("image", "")
|
||||
if not img_path:
|
||||
continue
|
||||
ext = Path(img_path).suffix.lower()
|
||||
extensions[ext] = extensions.get(ext, 0) + 1
|
||||
if Path(img_path).exists():
|
||||
existing += 1
|
||||
|
||||
vision_table = Table(title="Vision Stats")
|
||||
vision_table.add_column("Metric", style="bold")
|
||||
vision_table.add_column("Value")
|
||||
vision_table.add_row("Images referenced", str(has_image))
|
||||
vision_table.add_row("Missing image field", str(missing_image))
|
||||
vision_table.add_row("Images found on disk", str(existing))
|
||||
if extensions:
|
||||
ext_str = ", ".join(f"{ext} ({count})" for ext, count in sorted(extensions.items()))
|
||||
vision_table.add_row("Image formats", ext_str)
|
||||
console.print(vision_table)
|
||||
|
||||
|
||||
def _write_jsonl(path: Path, data: list[dict]) -> None:
|
||||
"""Write a list of dicts as JSONL."""
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ DEPS = [
|
|||
("deepspeed", "deepspeed", "0.12.0", False),
|
||||
("httpx", "httpx", "0.24.0", False),
|
||||
("unsloth", "unsloth", "2024.8", False),
|
||||
("PIL", "Pillow", "9.0.0", False),
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ def init(
|
|||
None,
|
||||
"--template",
|
||||
"-t",
|
||||
help="Use a template: chat, code, medical, reasoning",
|
||||
help="Use a template: chat, code, medical, reasoning, vision",
|
||||
),
|
||||
output: str = typer.Option(
|
||||
"soup.yaml",
|
||||
|
|
|
|||
|
|
@ -17,12 +17,16 @@ class LoraConfig(BaseModel):
|
|||
|
||||
class DataConfig(BaseModel):
|
||||
train: str = Field(..., description="Path to training data or HF dataset name")
|
||||
format: Literal["alpaca", "sharegpt", "chatml", "dpo", "auto"] = Field(
|
||||
format: Literal["alpaca", "sharegpt", "chatml", "dpo", "llava", "sharegpt4v", "auto"] = Field(
|
||||
default="auto",
|
||||
description="Data format",
|
||||
)
|
||||
val_split: float = Field(default=0.1, ge=0.0, le=0.5, description="Validation split ratio")
|
||||
max_length: int = Field(default=2048, description="Max sequence length in tokens")
|
||||
image_dir: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Base directory for resolving relative image paths in vision datasets",
|
||||
)
|
||||
|
||||
|
||||
class TrainingConfig(BaseModel):
|
||||
|
|
@ -67,6 +71,10 @@ class SoupConfig(BaseModel):
|
|||
|
||||
base: str = Field(..., description="Base model name or path (HF model ID)")
|
||||
task: Literal["sft", "dpo", "grpo"] = Field(default="sft", description="Training task type")
|
||||
modality: Literal["text", "vision"] = Field(
|
||||
default="text",
|
||||
description="Training modality: text (default) or vision (multimodal)",
|
||||
)
|
||||
backend: Literal["transformers", "unsloth"] = Field(
|
||||
default="transformers",
|
||||
description="Training backend: transformers (default) or unsloth (2-5x faster)",
|
||||
|
|
@ -157,6 +165,33 @@ training:
|
|||
num_generations: 4
|
||||
reward_fn: accuracy
|
||||
|
||||
output: ./output
|
||||
""",
|
||||
"vision": """# Soup template: Vision / Multimodal
|
||||
# Fine-tune a vision-language model for image understanding
|
||||
|
||||
base: meta-llama/Llama-3.2-11B-Vision-Instruct
|
||||
task: sft
|
||||
modality: vision
|
||||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||||
|
||||
data:
|
||||
train: ./data/vision_train.jsonl
|
||||
format: llava
|
||||
image_dir: ./data/images
|
||||
val_split: 0.1
|
||||
max_length: 2048
|
||||
|
||||
training:
|
||||
epochs: 3
|
||||
lr: 1e-5
|
||||
batch_size: auto
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
target_modules: auto
|
||||
quantization: 4bit
|
||||
|
||||
output: ./output
|
||||
""",
|
||||
"medical": """# Soup template: Medical / Domain Expert
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ Supported formats:
|
|||
- sharegpt: {"conversations": [{"from": "human", "value": ...}, ...]}
|
||||
- chatml: {"messages": [{"role": "user", "content": ...}, ...]}
|
||||
- dpo: {"prompt": ..., "chosen": ..., "rejected": ...}
|
||||
- llava: {"image": ..., "conversations": [{"from": "human", "value": ...}, ...]}
|
||||
- sharegpt4v: {"image": ..., "conversations": [{"from": "human", "value": ...}, ...]}
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
|
@ -19,6 +21,8 @@ FORMAT_SIGNATURES = {
|
|||
"sharegpt": {"conversations"},
|
||||
"chatml": {"messages"},
|
||||
"dpo": {"prompt", "chosen", "rejected"},
|
||||
"llava": {"image", "conversations"},
|
||||
"sharegpt4v": {"image", "conversations"},
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -30,7 +34,10 @@ def detect_format(data: list[dict]) -> str:
|
|||
sample = data[0]
|
||||
keys = set(sample.keys())
|
||||
|
||||
for fmt, required_keys in FORMAT_SIGNATURES.items():
|
||||
# Check more specific formats first (llava/sharegpt4v before sharegpt)
|
||||
check_order = ["alpaca", "llava", "dpo", "sharegpt", "chatml"]
|
||||
for fmt in check_order:
|
||||
required_keys = FORMAT_SIGNATURES[fmt]
|
||||
if required_keys.issubset(keys):
|
||||
return fmt
|
||||
|
||||
|
|
@ -38,7 +45,8 @@ def detect_format(data: list[dict]) -> str:
|
|||
f"Cannot detect format. Keys found: {keys}. "
|
||||
f"Expected one of: alpaca (instruction, output), "
|
||||
f"sharegpt (conversations), chatml (messages), "
|
||||
f"dpo (prompt, chosen, rejected)"
|
||||
f"dpo (prompt, chosen, rejected), "
|
||||
f"llava/sharegpt4v (image, conversations)"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -46,6 +54,7 @@ def format_to_messages(row: dict, fmt: str) -> Optional[dict]:
|
|||
"""Convert any format to unified messages format for training.
|
||||
|
||||
Returns: {"messages": [{"role": ..., "content": ...}, ...]}
|
||||
For vision formats, also includes "image" key.
|
||||
"""
|
||||
try:
|
||||
if fmt == "chatml":
|
||||
|
|
@ -56,6 +65,8 @@ def format_to_messages(row: dict, fmt: str) -> Optional[dict]:
|
|||
return _convert_sharegpt(row)
|
||||
elif fmt == "dpo":
|
||||
return _convert_dpo(row)
|
||||
elif fmt in ("llava", "sharegpt4v"):
|
||||
return _convert_vision(row)
|
||||
else:
|
||||
raise ValueError(f"Unknown format: {fmt}")
|
||||
except (KeyError, TypeError, IndexError):
|
||||
|
|
@ -106,6 +117,32 @@ def _convert_dpo(row: dict) -> dict:
|
|||
}
|
||||
|
||||
|
||||
def _convert_vision(row: dict) -> dict:
|
||||
"""Convert LLaVA / ShareGPT4V vision format to unified messages + image.
|
||||
|
||||
Input: {"image": "path.jpg", "conversations": [{"from": "human", "value": ...}, ...]}
|
||||
Output: {"messages": [...], "image": "path.jpg"}
|
||||
"""
|
||||
conversations = row["conversations"]
|
||||
role_map = {"human": "user", "gpt": "assistant", "system": "system"}
|
||||
|
||||
messages = []
|
||||
for turn in conversations:
|
||||
role = role_map.get(turn["from"], turn["from"])
|
||||
messages.append({"role": role, "content": turn["value"]})
|
||||
|
||||
result = {"messages": messages, "image": row["image"]}
|
||||
# Preserve optional id field
|
||||
if "id" in row:
|
||||
result["id"] = row["id"]
|
||||
return result
|
||||
|
||||
|
||||
def is_vision_format(fmt: str) -> bool:
|
||||
"""Check if a format is a vision/multimodal format."""
|
||||
return fmt in ("llava", "sharegpt4v")
|
||||
|
||||
|
||||
# --- Reverse conversion: messages → target format ---
|
||||
|
||||
CONVERTIBLE_FORMATS = ("alpaca", "sharegpt", "chatml")
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from pathlib import Path
|
|||
from rich.console import Console
|
||||
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.formats import detect_format, format_to_messages
|
||||
from soup_cli.data.formats import detect_format, format_to_messages, is_vision_format
|
||||
|
||||
console = Console()
|
||||
|
||||
|
|
@ -101,6 +101,11 @@ def load_dataset(data_config: DataConfig) -> dict:
|
|||
formatted = [format_to_messages(row, fmt) for row in raw_data]
|
||||
formatted = [r for r in formatted if r is not None] # filter failed rows
|
||||
|
||||
# Validate image paths for vision formats
|
||||
if is_vision_format(fmt):
|
||||
image_dir = Path(data_config.image_dir) if data_config.image_dir else path.parent
|
||||
formatted = _validate_vision_images(formatted, image_dir)
|
||||
|
||||
# Split into train/val
|
||||
if data_config.val_split > 0:
|
||||
split_idx = int(len(formatted) * (1 - data_config.val_split))
|
||||
|
|
@ -112,6 +117,29 @@ def load_dataset(data_config: DataConfig) -> dict:
|
|||
return {"train": formatted}
|
||||
|
||||
|
||||
def _validate_vision_images(data: list[dict], image_dir: Path) -> list[dict]:
|
||||
"""Validate and resolve image paths in vision dataset rows.
|
||||
|
||||
Each row must have an 'image' key with a filename or path.
|
||||
Resolves relative paths against image_dir.
|
||||
"""
|
||||
valid = []
|
||||
missing = 0
|
||||
for row in data:
|
||||
if "image" not in row or not row["image"]:
|
||||
missing += 1
|
||||
continue
|
||||
image_path = Path(row["image"])
|
||||
if not image_path.is_absolute():
|
||||
image_path = image_dir / image_path
|
||||
row["image"] = str(image_path)
|
||||
valid.append(row)
|
||||
|
||||
if missing > 0:
|
||||
console.print(f"[yellow]Warning: {missing} rows skipped (missing image path)[/]")
|
||||
return valid
|
||||
|
||||
|
||||
def _load_hf_dataset(name: str, data_config: DataConfig) -> dict:
|
||||
"""Load a dataset from HuggingFace Hub."""
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -42,8 +42,11 @@ class SFTTrainerWrapper:
|
|||
cfg = self.config
|
||||
tcfg = cfg.training
|
||||
use_unsloth = cfg.backend == "unsloth"
|
||||
use_vision = cfg.modality == "vision"
|
||||
|
||||
if use_unsloth:
|
||||
if use_vision:
|
||||
self._setup_vision_transformers(cfg, tcfg)
|
||||
elif use_unsloth:
|
||||
self._setup_unsloth(cfg, tcfg)
|
||||
else:
|
||||
self._setup_transformers(cfg, tcfg)
|
||||
|
|
@ -72,29 +75,32 @@ class SFTTrainerWrapper:
|
|||
console.print(f"[green]Auto batch size:[/] {batch_size}")
|
||||
|
||||
# --- Dataset ---
|
||||
def format_row(example):
|
||||
if hasattr(self.tokenizer, "chat_template") and self.tokenizer.chat_template:
|
||||
text = self.tokenizer.apply_chat_template(
|
||||
example["messages"], tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
else:
|
||||
# Fallback for models without chat template
|
||||
parts = []
|
||||
for msg in example["messages"]:
|
||||
role = msg["role"]
|
||||
content = msg["content"]
|
||||
parts.append(f"{role}: {content}")
|
||||
text = "\n".join(parts)
|
||||
return {"text": text}
|
||||
if use_vision:
|
||||
train_ds, eval_ds = self._prepare_vision_dataset(dataset)
|
||||
else:
|
||||
def format_row(example):
|
||||
if hasattr(self.tokenizer, "chat_template") and self.tokenizer.chat_template:
|
||||
text = self.tokenizer.apply_chat_template(
|
||||
example["messages"], tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
else:
|
||||
# Fallback for models without chat template
|
||||
parts = []
|
||||
for msg in example["messages"]:
|
||||
role = msg["role"]
|
||||
content = msg["content"]
|
||||
parts.append(f"{role}: {content}")
|
||||
text = "\n".join(parts)
|
||||
return {"text": text}
|
||||
|
||||
train_ds = Dataset.from_list(dataset["train"]).map(
|
||||
format_row, remove_columns=["messages"]
|
||||
)
|
||||
eval_ds = None
|
||||
if "val" in dataset and dataset["val"]:
|
||||
eval_ds = Dataset.from_list(dataset["val"]).map(
|
||||
train_ds = Dataset.from_list(dataset["train"]).map(
|
||||
format_row, remove_columns=["messages"]
|
||||
)
|
||||
eval_ds = None
|
||||
if "val" in dataset and dataset["val"]:
|
||||
eval_ds = Dataset.from_list(dataset["val"]).map(
|
||||
format_row, remove_columns=["messages"]
|
||||
)
|
||||
|
||||
# --- Output dir ---
|
||||
output_dir = Path(cfg.output)
|
||||
|
|
@ -209,6 +215,90 @@ class SFTTrainerWrapper:
|
|||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
def _setup_vision_transformers(self, cfg, tcfg):
|
||||
"""Load vision-language model via transformers (LLaMA-Vision, Qwen2-VL, etc.)."""
|
||||
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
||||
from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig
|
||||
|
||||
console.print(f"[dim]Loading vision processor: {cfg.base}[/]")
|
||||
self.processor = AutoProcessor.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
self.tokenizer = self.processor # SFTTrainer uses processing_class
|
||||
|
||||
# Quantization
|
||||
bnb_config = None
|
||||
if tcfg.quantization == "4bit":
|
||||
import torch
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
elif tcfg.quantization == "8bit":
|
||||
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
console.print(f"[dim]Loading vision model: {cfg.base}[/]")
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": "auto"}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
self.model = AutoModelForVision2Seq.from_pretrained(cfg.base, **model_kwargs)
|
||||
|
||||
if tcfg.quantization in ("4bit", "8bit"):
|
||||
self.model = prepare_model_for_kbit_training(self.model)
|
||||
|
||||
# LoRA — target language model layers only
|
||||
target_modules = tcfg.lora.target_modules
|
||||
if target_modules == "auto":
|
||||
target_modules = None
|
||||
|
||||
lora_config = LoraConfig(
|
||||
r=tcfg.lora.r,
|
||||
lora_alpha=tcfg.lora.alpha,
|
||||
lora_dropout=tcfg.lora.dropout,
|
||||
target_modules=target_modules,
|
||||
bias="none",
|
||||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
|
||||
def _prepare_vision_dataset(self, dataset: dict):
|
||||
"""Prepare dataset for vision fine-tuning with image loading."""
|
||||
from datasets import Dataset
|
||||
|
||||
def load_and_format_vision(example):
|
||||
from PIL import Image as PILImage
|
||||
|
||||
image_path = example.get("image", "")
|
||||
image = None
|
||||
if image_path:
|
||||
try:
|
||||
image = PILImage.open(image_path).convert("RGB")
|
||||
except (FileNotFoundError, OSError):
|
||||
console.print(f"[yellow]Warning: cannot open image: {image_path}[/]")
|
||||
|
||||
messages = example["messages"]
|
||||
text = self.processor.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
result = {"text": text}
|
||||
if image is not None:
|
||||
result["images"] = [image]
|
||||
return result
|
||||
|
||||
remove_cols = ["messages", "image"]
|
||||
train_ds = Dataset.from_list(dataset["train"]).map(
|
||||
load_and_format_vision,
|
||||
remove_columns=[c for c in remove_cols if c in dataset["train"][0]],
|
||||
)
|
||||
eval_ds = None
|
||||
if "val" in dataset and dataset["val"]:
|
||||
eval_ds = Dataset.from_list(dataset["val"]).map(
|
||||
load_and_format_vision,
|
||||
remove_columns=[c for c in remove_cols if c in dataset["val"][0]],
|
||||
)
|
||||
return train_ds, eval_ds
|
||||
|
||||
def train(
|
||||
self,
|
||||
display: Optional[object] = None,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,613 @@
|
|||
"""Tests for Phase 6 — Multimodal Fine-tuning (vision config, formats, loader, trainer)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from soup_cli.config.schema import TEMPLATES, SoupConfig
|
||||
from soup_cli.data.formats import (
|
||||
detect_format,
|
||||
format_to_messages,
|
||||
is_vision_format,
|
||||
)
|
||||
|
||||
# ─── Sample Data ───────────────────────────────────────────────────────────
|
||||
|
||||
LLAVA_ROW = {
|
||||
"image": "photo.jpg",
|
||||
"conversations": [
|
||||
{"from": "human", "value": "<image>\nDescribe this image."},
|
||||
{"from": "gpt", "value": "The image shows a cat sitting on a mat."},
|
||||
],
|
||||
}
|
||||
|
||||
SHAREGPT4V_ROW = {
|
||||
"image": "chart.png",
|
||||
"conversations": [
|
||||
{"from": "human", "value": "<image>\nWhat does this chart show?"},
|
||||
{"from": "gpt", "value": "The chart shows quarterly revenue growth."},
|
||||
],
|
||||
}
|
||||
|
||||
LLAVA_DATASET = [
|
||||
{
|
||||
"image": "img1.jpg",
|
||||
"conversations": [
|
||||
{"from": "human", "value": "What is this?"},
|
||||
{"from": "gpt", "value": "A dog."},
|
||||
],
|
||||
},
|
||||
{
|
||||
"image": "img2.jpg",
|
||||
"conversations": [
|
||||
{"from": "human", "value": "Describe the scene."},
|
||||
{"from": "gpt", "value": "A park with trees."},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
# ─── Config Tests ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestVisionConfig:
|
||||
"""Test modality config field validation."""
|
||||
|
||||
def test_modality_default_is_text(self):
|
||||
"""Default modality should be 'text'."""
|
||||
cfg = SoupConfig(base="some-model", data={"train": "./data.jsonl"})
|
||||
assert cfg.modality == "text"
|
||||
|
||||
def test_modality_vision_accepted(self):
|
||||
"""modality: vision should be valid."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
modality="vision",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.modality == "vision"
|
||||
|
||||
def test_modality_text_accepted(self):
|
||||
"""modality: text should be valid."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
modality="text",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.modality == "text"
|
||||
|
||||
def test_modality_invalid_rejected(self):
|
||||
"""Invalid modality should raise validation error."""
|
||||
with pytest.raises(Exception):
|
||||
SoupConfig(
|
||||
base="some-model",
|
||||
modality="audio",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
|
||||
def test_modality_in_model_dump(self):
|
||||
"""modality field should appear in model_dump output."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
modality="vision",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
dump = cfg.model_dump()
|
||||
assert dump["modality"] == "vision"
|
||||
|
||||
def test_vision_with_sft(self):
|
||||
"""Vision modality should work with SFT task."""
|
||||
cfg = SoupConfig(
|
||||
base="meta-llama/Llama-3.2-11B-Vision-Instruct",
|
||||
task="sft",
|
||||
modality="vision",
|
||||
data={"train": "./data.jsonl", "format": "llava"},
|
||||
)
|
||||
assert cfg.task == "sft"
|
||||
assert cfg.modality == "vision"
|
||||
|
||||
def test_full_vision_config(self):
|
||||
"""Full config with vision modality should validate."""
|
||||
cfg = SoupConfig(
|
||||
base="meta-llama/Llama-3.2-11B-Vision-Instruct",
|
||||
task="sft",
|
||||
modality="vision",
|
||||
data={
|
||||
"train": "./data.jsonl",
|
||||
"format": "llava",
|
||||
"image_dir": "./images",
|
||||
"max_length": 2048,
|
||||
},
|
||||
training={
|
||||
"epochs": 3,
|
||||
"lr": 1e-5,
|
||||
"quantization": "4bit",
|
||||
"lora": {"r": 64, "alpha": 16},
|
||||
},
|
||||
)
|
||||
assert cfg.modality == "vision"
|
||||
assert cfg.data.format == "llava"
|
||||
assert cfg.data.image_dir == "./images"
|
||||
|
||||
def test_vision_with_backend(self):
|
||||
"""Vision modality should work with different backends."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
modality="vision",
|
||||
backend="transformers",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.modality == "vision"
|
||||
assert cfg.backend == "transformers"
|
||||
|
||||
|
||||
# ─── Data Format Config Tests ─────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestVisionDataConfig:
|
||||
"""Test vision-related fields in DataConfig."""
|
||||
|
||||
def test_image_dir_default_none(self):
|
||||
"""image_dir should default to None."""
|
||||
cfg = SoupConfig(base="some-model", data={"train": "./data.jsonl"})
|
||||
assert cfg.data.image_dir is None
|
||||
|
||||
def test_image_dir_accepted(self):
|
||||
"""image_dir should be accepted as a string."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
data={"train": "./data.jsonl", "image_dir": "./images"},
|
||||
)
|
||||
assert cfg.data.image_dir == "./images"
|
||||
|
||||
def test_llava_format_accepted(self):
|
||||
"""format: llava should be valid."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
data={"train": "./data.jsonl", "format": "llava"},
|
||||
)
|
||||
assert cfg.data.format == "llava"
|
||||
|
||||
def test_sharegpt4v_format_accepted(self):
|
||||
"""format: sharegpt4v should be valid."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
data={"train": "./data.jsonl", "format": "sharegpt4v"},
|
||||
)
|
||||
assert cfg.data.format == "sharegpt4v"
|
||||
|
||||
|
||||
# ─── Format Detection Tests ───────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestVisionFormatDetection:
|
||||
"""Test auto-detection of vision formats."""
|
||||
|
||||
def test_detect_llava_format(self):
|
||||
"""Should detect LLaVA format (image + conversations)."""
|
||||
result = detect_format(LLAVA_DATASET)
|
||||
assert result == "llava"
|
||||
|
||||
def test_detect_sharegpt4v_with_explicit_format(self):
|
||||
"""ShareGPT4V has same structure as LLaVA; auto-detect returns llava."""
|
||||
data = [SHAREGPT4V_ROW]
|
||||
result = detect_format(data)
|
||||
# Both have same keys, so llava is detected first
|
||||
assert result == "llava"
|
||||
|
||||
def test_detect_llava_not_confused_with_sharegpt(self):
|
||||
"""LLaVA format (image + conversations) should not be detected as sharegpt."""
|
||||
result = detect_format([LLAVA_ROW])
|
||||
assert result != "sharegpt"
|
||||
|
||||
def test_detect_sharegpt_without_image(self):
|
||||
"""Regular ShareGPT (no image key) should still be detected as sharegpt."""
|
||||
data = [{"conversations": [{"from": "human", "value": "Hi"}]}]
|
||||
result = detect_format(data)
|
||||
assert result == "sharegpt"
|
||||
|
||||
def test_is_vision_format_llava(self):
|
||||
"""is_vision_format should return True for llava."""
|
||||
assert is_vision_format("llava") is True
|
||||
|
||||
def test_is_vision_format_sharegpt4v(self):
|
||||
"""is_vision_format should return True for sharegpt4v."""
|
||||
assert is_vision_format("sharegpt4v") is True
|
||||
|
||||
def test_is_vision_format_alpaca(self):
|
||||
"""is_vision_format should return False for alpaca."""
|
||||
assert is_vision_format("alpaca") is False
|
||||
|
||||
def test_is_vision_format_sharegpt(self):
|
||||
"""is_vision_format should return False for sharegpt."""
|
||||
assert is_vision_format("sharegpt") is False
|
||||
|
||||
|
||||
# ─── Format Conversion Tests ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestVisionFormatConversion:
|
||||
"""Test conversion of vision formats to unified message format."""
|
||||
|
||||
def test_llava_to_messages(self):
|
||||
"""LLaVA row should convert to messages + image."""
|
||||
result = format_to_messages(LLAVA_ROW, "llava")
|
||||
assert result is not None
|
||||
assert "messages" in result
|
||||
assert "image" in result
|
||||
assert result["image"] == "photo.jpg"
|
||||
assert len(result["messages"]) == 2
|
||||
assert result["messages"][0]["role"] == "user"
|
||||
assert result["messages"][1]["role"] == "assistant"
|
||||
|
||||
def test_sharegpt4v_to_messages(self):
|
||||
"""ShareGPT4V row should convert to messages + image."""
|
||||
result = format_to_messages(SHAREGPT4V_ROW, "sharegpt4v")
|
||||
assert result is not None
|
||||
assert "messages" in result
|
||||
assert result["image"] == "chart.png"
|
||||
assert result["messages"][0]["content"] == "<image>\nWhat does this chart show?"
|
||||
|
||||
def test_llava_preserves_image_tag(self):
|
||||
"""<image> tag in conversation should be preserved."""
|
||||
result = format_to_messages(LLAVA_ROW, "llava")
|
||||
assert "<image>" in result["messages"][0]["content"]
|
||||
|
||||
def test_llava_with_id(self):
|
||||
"""LLaVA row with id field should preserve it."""
|
||||
row = {**LLAVA_ROW, "id": "sample_001"}
|
||||
result = format_to_messages(row, "llava")
|
||||
assert result["id"] == "sample_001"
|
||||
|
||||
def test_llava_role_mapping(self):
|
||||
"""human/gpt roles should map to user/assistant."""
|
||||
row = {
|
||||
"image": "test.jpg",
|
||||
"conversations": [
|
||||
{"from": "system", "value": "You are helpful."},
|
||||
{"from": "human", "value": "Describe this."},
|
||||
{"from": "gpt", "value": "A test image."},
|
||||
],
|
||||
}
|
||||
result = format_to_messages(row, "llava")
|
||||
roles = [msg["role"] for msg in result["messages"]]
|
||||
assert roles == ["system", "user", "assistant"]
|
||||
|
||||
def test_llava_invalid_row_returns_none(self):
|
||||
"""Row missing required keys should return None."""
|
||||
result = format_to_messages({"bad": "data"}, "llava")
|
||||
assert result is None
|
||||
|
||||
def test_sharegpt4v_invalid_row_returns_none(self):
|
||||
"""Row missing required keys should return None."""
|
||||
result = format_to_messages({"bad": "data"}, "sharegpt4v")
|
||||
assert result is None
|
||||
|
||||
|
||||
# ─── Data Loader Vision Tests ─────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestVisionDataLoader:
|
||||
"""Test vision image validation in data loader."""
|
||||
|
||||
def test_validate_vision_images_resolves_relative(self):
|
||||
"""Relative image paths should be resolved against image_dir."""
|
||||
from soup_cli.data.loader import _validate_vision_images
|
||||
|
||||
data = [{"messages": [{"role": "user", "content": "Hi"}], "image": "photo.jpg"}]
|
||||
image_dir = Path("/data/images")
|
||||
result = _validate_vision_images(data, image_dir)
|
||||
assert len(result) == 1
|
||||
assert result[0]["image"] == str(image_dir / "photo.jpg")
|
||||
|
||||
def test_validate_vision_images_skips_missing(self):
|
||||
"""Rows without image field should be skipped."""
|
||||
from soup_cli.data.loader import _validate_vision_images
|
||||
|
||||
data = [
|
||||
{"messages": [{"role": "user", "content": "Hi"}], "image": "photo.jpg"},
|
||||
{"messages": [{"role": "user", "content": "No image"}]},
|
||||
]
|
||||
image_dir = Path("/data/images")
|
||||
result = _validate_vision_images(data, image_dir)
|
||||
assert len(result) == 1
|
||||
|
||||
def test_validate_vision_images_skips_empty_image(self):
|
||||
"""Rows with empty image string should be skipped."""
|
||||
from soup_cli.data.loader import _validate_vision_images
|
||||
|
||||
data = [{"messages": [{"role": "user", "content": "Hi"}], "image": ""}]
|
||||
image_dir = Path("/data")
|
||||
result = _validate_vision_images(data, image_dir)
|
||||
assert len(result) == 0
|
||||
|
||||
def test_validate_vision_images_absolute_path(self):
|
||||
"""Absolute image paths should be kept as-is."""
|
||||
from soup_cli.data.loader import _validate_vision_images
|
||||
|
||||
abs_path = str(Path("/absolute/path/photo.jpg"))
|
||||
data = [{"messages": [{"role": "user", "content": "Hi"}], "image": abs_path}]
|
||||
image_dir = Path("/different/dir")
|
||||
result = _validate_vision_images(data, image_dir)
|
||||
assert result[0]["image"] == abs_path
|
||||
|
||||
def test_load_dataset_with_vision_format(self):
|
||||
"""load_dataset should handle llava format data files."""
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.loader import load_dataset
|
||||
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w", suffix=".jsonl", delete=False, encoding="utf-8"
|
||||
) as tmpfile:
|
||||
for row in LLAVA_DATASET:
|
||||
tmpfile.write(json.dumps(row) + "\n")
|
||||
tmpfile_path = tmpfile.name
|
||||
|
||||
try:
|
||||
data_config = DataConfig(
|
||||
train=tmpfile_path,
|
||||
format="llava",
|
||||
val_split=0.0,
|
||||
)
|
||||
result = load_dataset(data_config)
|
||||
assert "train" in result
|
||||
assert len(result["train"]) == 2
|
||||
# Each row should have messages and image
|
||||
for row in result["train"]:
|
||||
assert "messages" in row
|
||||
assert "image" in row
|
||||
finally:
|
||||
Path(tmpfile_path).unlink()
|
||||
|
||||
|
||||
# ─── SFT Trainer Vision Tests ─────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestSFTVisionIntegration:
|
||||
"""Test SFT trainer with vision modality."""
|
||||
|
||||
def test_sft_wrapper_init_with_vision(self):
|
||||
"""SFTTrainerWrapper should accept vision modality config."""
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="meta-llama/Llama-3.2-11B-Vision-Instruct",
|
||||
modality="vision",
|
||||
data={"train": "./data.jsonl", "format": "llava"},
|
||||
)
|
||||
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
||||
assert wrapper.config.modality == "vision"
|
||||
|
||||
def test_sft_setup_vision_calls_automodel(self):
|
||||
"""_setup_vision_transformers should use AutoModelForVision2Seq."""
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="test-vision-model",
|
||||
modality="vision",
|
||||
data={"train": "./data.jsonl", "format": "llava", "max_length": 2048},
|
||||
training={"lora": {"r": 64, "alpha": 16, "dropout": 0.05}},
|
||||
)
|
||||
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
||||
|
||||
mock_model = MagicMock()
|
||||
mock_model.get_nb_trainable_parameters.return_value = (1000, 100000)
|
||||
mock_processor = MagicMock()
|
||||
|
||||
with patch(
|
||||
"soup_cli.trainer.sft.SFTTrainerWrapper._setup_vision_transformers"
|
||||
) as mock_setup:
|
||||
mock_setup.side_effect = lambda c, t: setattr(wrapper, "model", mock_model) or setattr(
|
||||
wrapper, "tokenizer", mock_processor
|
||||
) or setattr(wrapper, "processor", mock_processor)
|
||||
wrapper._setup_vision_transformers(cfg, cfg.training)
|
||||
mock_setup.assert_called_once()
|
||||
|
||||
def test_sft_vision_config_selects_vision_path(self):
|
||||
"""Vision modality config should be detected correctly in wrapper."""
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="test-vision-model",
|
||||
modality="vision",
|
||||
data={"train": "./data.jsonl", "format": "llava"},
|
||||
)
|
||||
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
||||
# Verify config is stored correctly for vision routing
|
||||
assert wrapper.config.modality == "vision"
|
||||
assert wrapper.config.data.format == "llava"
|
||||
|
||||
def test_sft_vision_has_setup_methods(self):
|
||||
"""SFTTrainerWrapper should have vision-specific methods."""
|
||||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="test-vision-model",
|
||||
modality="vision",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
wrapper = SFTTrainerWrapper(cfg, device="cuda")
|
||||
assert hasattr(wrapper, "_setup_vision_transformers")
|
||||
assert hasattr(wrapper, "_prepare_vision_dataset")
|
||||
|
||||
|
||||
# ─── Template Tests ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestVisionTemplate:
|
||||
"""Test vision template in TEMPLATES dict."""
|
||||
|
||||
def test_vision_template_exists(self):
|
||||
"""Vision template should exist."""
|
||||
assert "vision" in TEMPLATES
|
||||
|
||||
def test_vision_template_has_modality(self):
|
||||
"""Vision template should set modality: vision."""
|
||||
assert "modality: vision" in TEMPLATES["vision"]
|
||||
|
||||
def test_vision_template_has_llava_format(self):
|
||||
"""Vision template should use llava format."""
|
||||
assert "format: llava" in TEMPLATES["vision"]
|
||||
|
||||
def test_vision_template_has_image_dir(self):
|
||||
"""Vision template should have image_dir field."""
|
||||
assert "image_dir:" in TEMPLATES["vision"]
|
||||
|
||||
def test_vision_template_has_vision_model(self):
|
||||
"""Vision template should use a vision model."""
|
||||
assert "Vision" in TEMPLATES["vision"]
|
||||
|
||||
def test_vision_template_mentions_unsloth(self):
|
||||
"""Vision template should mention unsloth as an option."""
|
||||
assert "unsloth" in TEMPLATES["vision"]
|
||||
|
||||
def test_vision_template_is_valid_yaml(self):
|
||||
"""Vision template should be valid YAML that parses."""
|
||||
import yaml
|
||||
|
||||
config = yaml.safe_load(TEMPLATES["vision"])
|
||||
assert config["modality"] == "vision"
|
||||
assert config["data"]["format"] == "llava"
|
||||
assert config["data"]["image_dir"] == "./data/images"
|
||||
|
||||
def test_vision_template_validates_as_config(self):
|
||||
"""Vision template should validate as a SoupConfig."""
|
||||
import yaml
|
||||
|
||||
config_dict = yaml.safe_load(TEMPLATES["vision"])
|
||||
cfg = SoupConfig(**config_dict)
|
||||
assert cfg.modality == "vision"
|
||||
assert cfg.data.format == "llava"
|
||||
assert cfg.data.image_dir == "./data/images"
|
||||
|
||||
|
||||
# ─── Init Command Tests ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestInitVisionTemplate:
|
||||
"""Test init command with --template vision."""
|
||||
|
||||
def test_init_vision_template(self):
|
||||
"""soup init --template vision should create a valid config."""
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
output_path = Path(tmpdir) / "soup.yaml"
|
||||
result = runner.invoke(app, ["init", "--template", "vision", "-o", str(output_path)])
|
||||
assert result.exit_code == 0
|
||||
assert output_path.exists()
|
||||
content = output_path.read_text()
|
||||
assert "modality: vision" in content
|
||||
assert "format: llava" in content
|
||||
|
||||
|
||||
# ─── Data Inspect Vision Tests ─────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestDataInspectVision:
|
||||
"""Test data inspect command with vision datasets."""
|
||||
|
||||
def test_show_vision_stats_with_images(self):
|
||||
"""_show_vision_stats should detect image fields and show stats."""
|
||||
from io import StringIO
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from soup_cli.commands.data import _show_vision_stats
|
||||
|
||||
data = [
|
||||
{"image": "photo1.jpg", "conversations": []},
|
||||
{"image": "photo2.png", "conversations": []},
|
||||
{"image": "", "conversations": []},
|
||||
]
|
||||
|
||||
output = StringIO()
|
||||
with patch("soup_cli.commands.data.console", Console(file=output)):
|
||||
_show_vision_stats(data)
|
||||
|
||||
text = output.getvalue()
|
||||
assert "Vision Stats" in text
|
||||
assert "2" in text # 2 images referenced
|
||||
|
||||
def test_show_vision_stats_no_images(self):
|
||||
"""_show_vision_stats should not print anything for non-vision datasets."""
|
||||
from io import StringIO
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from soup_cli.commands.data import _show_vision_stats
|
||||
|
||||
data = [{"instruction": "Hi", "output": "Hello"}]
|
||||
|
||||
output = StringIO()
|
||||
with patch("soup_cli.commands.data.console", Console(file=output)):
|
||||
_show_vision_stats(data)
|
||||
|
||||
text = output.getvalue()
|
||||
assert "Vision" not in text
|
||||
|
||||
def test_show_vision_stats_empty_data(self):
|
||||
"""_show_vision_stats should handle empty data gracefully."""
|
||||
from soup_cli.commands.data import _show_vision_stats
|
||||
|
||||
# Should not raise
|
||||
_show_vision_stats([])
|
||||
|
||||
def test_show_vision_stats_image_extensions(self):
|
||||
"""_show_vision_stats should report image file extensions."""
|
||||
from io import StringIO
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from soup_cli.commands.data import _show_vision_stats
|
||||
|
||||
data = [
|
||||
{"image": "a.jpg", "conversations": []},
|
||||
{"image": "b.jpg", "conversations": []},
|
||||
{"image": "c.png", "conversations": []},
|
||||
]
|
||||
|
||||
output = StringIO()
|
||||
with patch("soup_cli.commands.data.console", Console(file=output)):
|
||||
_show_vision_stats(data)
|
||||
|
||||
text = output.getvalue()
|
||||
assert ".jpg" in text
|
||||
assert ".png" in text
|
||||
|
||||
|
||||
# ─── Doctor Tests ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestDoctorVision:
|
||||
"""Test that doctor checks for Pillow (vision dependency)."""
|
||||
|
||||
def test_pillow_in_deps_list(self):
|
||||
"""Pillow should be listed in doctor dependencies."""
|
||||
from soup_cli.commands.doctor import DEPS
|
||||
|
||||
pkg_names = [pkg_name for _, pkg_name, _, _ in DEPS]
|
||||
assert "Pillow" in pkg_names or "pillow" in [n.lower() for n in pkg_names]
|
||||
|
||||
|
||||
# ─── Sweep Shortcut Tests ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestModalitySweepParam:
|
||||
"""Test modality parameter in sweep shortcuts."""
|
||||
|
||||
def test_modality_shortcut(self):
|
||||
"""modality should be settable via sweep param."""
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {"modality": "text"}
|
||||
_set_nested_param(config, "modality", "vision")
|
||||
assert config["modality"] == "vision"
|
||||
Loading…
Reference in New Issue