| """ |
| MINDI 1.5 Vision-Coder — Code Generation Model |
| |
| Loads the base coding model with LoRA adapters for fine-tuning |
| on Next.js + Tailwind + TypeScript code generation. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Optional |
|
|
| import torch |
| from peft import LoraConfig, PeftModel, get_peft_model, TaskType |
| from transformers import AutoModelForCausalLM, BitsAndBytesConfig |
|
|
|
|
| class MindiCodeModel: |
| """Base coding model with LoRA for MINDI 1.5 fine-tuning.""" |
|
|
| def __init__( |
| self, |
| model_name: str = "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct", |
| device: Optional[str] = None, |
| cache_dir: Optional[Path] = None, |
| load_in_4bit: bool = False, |
| ) -> None: |
| self.model_name = model_name |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| self.cache_dir = cache_dir or Path("./checkpoints/base") |
| self.cache_dir.mkdir(parents=True, exist_ok=True) |
| self.load_in_4bit = load_in_4bit |
| self.model: Optional[AutoModelForCausalLM] = None |
| self.peft_model: Optional[PeftModel] = None |
|
|
| def load_base_model(self) -> AutoModelForCausalLM: |
| """Load the base model with optional 4-bit quantization.""" |
| quantization_config = None |
| if self.load_in_4bit: |
| quantization_config = BitsAndBytesConfig( |
| load_in_4bit=True, |
| bnb_4bit_quant_type="nf4", |
| bnb_4bit_compute_dtype=torch.bfloat16, |
| bnb_4bit_use_double_quant=True, |
| ) |
|
|
| self.model = AutoModelForCausalLM.from_pretrained( |
| self.model_name, |
| cache_dir=str(self.cache_dir), |
| torch_dtype=torch.bfloat16, |
| device_map="auto" if self.device == "cuda" else None, |
| quantization_config=quantization_config, |
| trust_remote_code=True, |
| ) |
| return self.model |
|
|
| def apply_lora( |
| self, |
| rank: int = 64, |
| alpha: int = 128, |
| dropout: float = 0.05, |
| target_modules: Optional[list[str]] = None, |
| ) -> PeftModel: |
| """Apply LoRA adapters to the base model for efficient fine-tuning.""" |
| if self.model is None: |
| raise RuntimeError("Base model not loaded. Call load_base_model() first.") |
|
|
| if target_modules is None: |
| target_modules = [ |
| "q_proj", "k_proj", "v_proj", "o_proj", |
| "gate_proj", "up_proj", "down_proj", |
| ] |
|
|
| lora_config = LoraConfig( |
| r=rank, |
| lora_alpha=alpha, |
| lora_dropout=dropout, |
| target_modules=target_modules, |
| bias="none", |
| task_type=TaskType.CAUSAL_LM, |
| ) |
|
|
| self.peft_model = get_peft_model(self.model, lora_config) |
| trainable, total = self._count_parameters() |
| print(f"[MindiCodeModel] LoRA applied — trainable: {trainable:,} / {total:,} " |
| f"({100 * trainable / total:.2f}%)") |
| return self.peft_model |
|
|
| def _count_parameters(self) -> tuple[int, int]: |
| """Count trainable and total parameters.""" |
| model = self.peft_model or self.model |
| if model is None: |
| return 0, 0 |
| trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| total = sum(p.numel() for p in model.parameters()) |
| return trainable, total |
|
|
| def save_adapter(self, output_dir: Optional[Path] = None) -> Path: |
| """Save the LoRA adapter weights.""" |
| if self.peft_model is None: |
| raise RuntimeError("No LoRA adapter to save. Call apply_lora() first.") |
| save_path = output_dir or Path("./checkpoints/finetuned") |
| save_path.mkdir(parents=True, exist_ok=True) |
| self.peft_model.save_pretrained(str(save_path)) |
| return save_path |
|
|
| def load_adapter(self, adapter_dir: Path) -> PeftModel: |
| """Load a saved LoRA adapter onto the base model.""" |
| if self.model is None: |
| self.load_base_model() |
| self.peft_model = PeftModel.from_pretrained( |
| self.model, str(adapter_dir) |
| ) |
| return self.peft_model |
|
|
| def resize_embeddings(self, new_vocab_size: int) -> None: |
| """Resize model embeddings to accommodate new special tokens.""" |
| model = self.peft_model or self.model |
| if model is None: |
| raise RuntimeError("No model loaded.") |
| model.resize_token_embeddings(new_vocab_size) |
|
|
