From 5b7ad3c35869dcad8b6e5943a4a9e20e4266a01e Mon Sep 17 00:00:00 2001 From: Alpamys Date: Mon, 23 Mar 2026 20:16:24 +0500 Subject: [PATCH] =?UTF-8?q?Add=20multimodal=20vision=20fine-tuning=20suppo?= =?UTF-8?q?rt=20(Phase=206)=20=E2=80=94=20v0.5.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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) --- CLAUDE.md | 7 +- README.md | 63 +++- pyproject.toml | 3 +- soup_cli/__init__.py | 2 +- soup_cli/commands/data.py | 41 +++ soup_cli/commands/doctor.py | 1 + soup_cli/commands/init.py | 2 +- soup_cli/config/schema.py | 37 ++- soup_cli/data/formats.py | 41 ++- soup_cli/data/loader.py | 30 +- soup_cli/trainer/sft.py | 132 ++++++-- tests/test_vision.py | 613 ++++++++++++++++++++++++++++++++++++ 12 files changed, 940 insertions(+), 32 deletions(-) create mode 100644 tests/test_vision.py diff --git a/CLAUDE.md b/CLAUDE.md index 769c174..dc88156 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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 | diff --git a/README.md b/README.md index c93fdf2..eaf5279 100644 --- a/README.md +++ b/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": "\nDescribe this image."}, {"from": "gpt", "value": "A cat on a mat."}]} +``` + +**ShareGPT4V:** +```json +{"image": "chart.png", "conversations": [{"from": "human", "value": "\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": "\nDescribe this."}, {"from": "gpt", "value": "A cat."}]} +``` + +**ShareGPT4V (vision):** +```json +{"image": "chart.png", "conversations": [{"from": "human", "value": "\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 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) | diff --git a/pyproject.toml b/pyproject.toml index 48ffd3c..f8eca4f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/soup_cli/__init__.py b/soup_cli/__init__.py index 1311874..949edfe 100644 --- a/soup_cli/__init__.py +++ b/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune LLMs in one command.""" -__version__ = "0.4.3" +__version__ = "0.5.0" diff --git a/soup_cli/commands/data.py b/soup_cli/commands/data.py index 09b7d74..853e40d 100644 --- a/soup_cli/commands/data.py +++ b/soup_cli/commands/data.py @@ -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: diff --git a/soup_cli/commands/doctor.py b/soup_cli/commands/doctor.py index 6ab08e1..6fdd070 100644 --- a/soup_cli/commands/doctor.py +++ b/soup_cli/commands/doctor.py @@ -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), ] diff --git a/soup_cli/commands/init.py b/soup_cli/commands/init.py index d586595..9ae17f2 100644 --- a/soup_cli/commands/init.py +++ b/soup_cli/commands/init.py @@ -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", diff --git a/soup_cli/config/schema.py b/soup_cli/config/schema.py index abc0788..9c8bc3c 100644 --- a/soup_cli/config/schema.py +++ b/soup_cli/config/schema.py @@ -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 diff --git a/soup_cli/data/formats.py b/soup_cli/data/formats.py index 1136eee..255ed4c 100644 --- a/soup_cli/data/formats.py +++ b/soup_cli/data/formats.py @@ -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") diff --git a/soup_cli/data/loader.py b/soup_cli/data/loader.py index 4f7bd40..cb6f4b7 100644 --- a/soup_cli/data/loader.py +++ b/soup_cli/data/loader.py @@ -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: diff --git a/soup_cli/trainer/sft.py b/soup_cli/trainer/sft.py index c5179e3..d2c49cd 100644 --- a/soup_cli/trainer/sft.py +++ b/soup_cli/trainer/sft.py @@ -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, diff --git a/tests/test_vision.py b/tests/test_vision.py new file mode 100644 index 0000000..37e8183 --- /dev/null +++ b/tests/test_vision.py @@ -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": "\nDescribe this image."}, + {"from": "gpt", "value": "The image shows a cat sitting on a mat."}, + ], +} + +SHAREGPT4V_ROW = { + "image": "chart.png", + "conversations": [ + {"from": "human", "value": "\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"] == "\nWhat does this chart show?" + + def test_llava_preserves_image_tag(self): + """ tag in conversation should be preserved.""" + result = format_to_messages(LLAVA_ROW, "llava") + assert "" 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"