Add smoke_test.py for local environment validation

smoke_test.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""
+smoke_test.py — 本機快速驗證
+執行時間 < 2 min；確認 deepseek_vl 可載入、processor 可處理 ChartQA
+"""
+import sys, subprocess, logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
+log = logging.getLogger(__name__)
+
+# ─── Install deepseek_vl if missing ──────────────────────────────────────────
+try:
+    from deepseek_vl.models import DeepseekVLV2Processor
+except ImportError:
+    log.info("Installing deepseek_vl …")
+    subprocess.run(
+        [sys.executable, "-m", "pip", "install", "-q",
+         "git+https://github.com/deepseek-ai/DeepSeek-VL2.git"],
+        check=True,
+    )
+    from deepseek_vl.models import DeepseekVLV2Processor
+
+import torch
+from datasets import load_dataset
+from PIL import Image
+from transformers import AutoModelForCausalLM
+from peft import LoraConfig, get_peft_model, TaskType
+
+MODEL_ID = "deepseek-ai/deepseek-vl2-tiny"
+
+# ─── 1. Processor ────────────────────────────────────────────────────────────
+log.info("Loading processor …")
+proc = DeepseekVLV2Processor.from_pretrained(MODEL_ID)
+log.info("Processor OK ✓")
+
+# ─── 2. ChartQA mini sample ──────────────────────────────────────────────────
+log.info("Loading 4 ChartQA samples …")
+ds = load_dataset("HuggingFaceM4/ChartQA", split="val[:4]")
+for row in ds:
+    img = row["image"]
+    if not isinstance(img, Image.Image):
+        img = Image.fromarray(img)
+    img  = img.convert("RGB")
+    q    = str(row["query"])
+    ans  = row["label"][0] if isinstance(row["label"], list) else str(row["label"])
+    conv = [
+        {"role": "<|User|>",      "content": f"<image>\n{q}", "images": [img]},
+        {"role": "<|Assistant|>", "content": ans},
+    ]
+    out = proc(conversations=[conv], images=[img], force_batchify=True, system_prompt="")
+    log.info(f"  input_ids shape = {out['input_ids'].shape}  query='{q[:40]}...'")
+log.info("Processor + ChartQA collation OK ✓")
+
+# ─── 3. Model load + LoRA (no forward pass — saves time) ─────────────────────
+log.info("Loading model (this takes ~1–2 min on first run) …")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, trust_remote_code=True, torch_dtype=torch.bfloat16,
+)
+lora = LoraConfig(
+    task_type=TaskType.CAUSAL_LM, r=16, lora_alpha=32,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+    bias="none",
+)
+model = get_peft_model(model, lora)
+model.print_trainable_parameters()
+log.info("LoRA wrapping OK ✓")
+
+if torch.cuda.is_available():
+    model = model.to("cuda")
+    mem = torch.cuda.memory_reserved() / 1e9
+    log.info(f"VRAM reserved = {mem:.1f} GB")
+    if mem > 11.5:
+        log.warning("VRAM > 11.5 GB — training with batch=1 might OOM. "
+                    "Try reducing MAX_TRAIN or set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True")
+    else:
+        log.info("VRAM looks fine for batch_size=1 training ✓")
+
+log.info("=" * 50)
+log.info("Smoke test PASSED — you can now run: python train_pipeline.py")
+log.info("=" * 50)