Spaces:
Sleeping
Sleeping
| import os | |
| import logging | |
| import torch | |
| from pathlib import Path | |
| from tqdm import tqdm | |
| from peft import LoraConfig, get_peft_model, PeftModel | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from torch.utils.data import DataLoader, Dataset | |
| os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0" | |
| logger = logging.getLogger(__name__) | |
| class SpiderDataset(Dataset): | |
| """Formats Spider JSON data into causal LM training examples (Question -> MongoDB Query).""" | |
| def __init__(self, data, tokenizer, max_length=512): | |
| self.data = data | |
| self.tokenizer = tokenizer | |
| self.max_length = max_length | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, idx): | |
| item = self.data[idx] | |
| question = item.get("question", "") | |
| mongodb_query = item.get("generated_mongodb_query", "") | |
| # Handle if stored as a list | |
| if isinstance(mongodb_query, list) and len(mongodb_query) > 0: | |
| mongodb_query = mongodb_query[0] | |
| prompt = f"### Task: Generate MongoDB Query (MQL)\n### Question:\n{question}\n\n### MongoDB Query:\n" | |
| prompt_enc = self.tokenizer(prompt, truncation=True, max_length=self.max_length, return_tensors="pt") | |
| completion_enc = self.tokenizer(mongodb_query, truncation=True, max_length=self.max_length, return_tensors="pt") | |
| input_ids = torch.cat([prompt_enc.input_ids[0], completion_enc.input_ids[0]], dim=0) | |
| attention_mask = torch.cat([prompt_enc.attention_mask[0], completion_enc.attention_mask[0]], dim=0) | |
| prompt_len = prompt_enc.input_ids.shape[1] | |
| labels = torch.cat([torch.full((prompt_len,), -100, dtype=torch.long), completion_enc.input_ids[0]], dim=0) | |
| return { | |
| "input_ids": input_ids, | |
| "attention_mask": attention_mask, | |
| "labels": labels | |
| } | |
| class ModelLoader: | |
| def __init__(self, config, device_override=None): | |
| self.config = config | |
| self.device = device_override or self._detect_device() | |
| if self.device == "cuda": | |
| self.dtype = torch.float16 | |
| elif self.device == "mps": | |
| self.dtype = torch.float16 | |
| else: | |
| self.dtype = torch.float32 | |
| self.tokenizer = None | |
| self.base_model = None | |
| self.lora_model = None | |
| logger.info(f"🖥️ Initialized ModelLoader | Device: {self.device} | Dtype: {self.dtype}") | |
| def _detect_device(self): | |
| if torch.cuda.is_available(): | |
| return "cuda" | |
| if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| return "mps" | |
| return "cpu" | |
| def load_tokenizer(self): | |
| model_name = getattr(self.config, "model_name", "Qwen/Qwen2.5-Coder-0.5B-Instruct") | |
| logger.info(f"🔄 Loading tokenizer: {model_name}") | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) | |
| if self.tokenizer.pad_token is None: | |
| self.tokenizer.pad_token = self.tokenizer.eos_token | |
| return self.tokenizer | |
| def load_base_model(self): | |
| model_name = getattr(self.config, "model_name", "Qwen/Qwen2.5-Coder-0.5B-Instruct") | |
| logger.info(f"🔄 Loading base model: {model_name}") | |
| if self.device == "cuda": | |
| self.base_model = AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| torch_dtype=self.dtype, | |
| trust_remote_code=True, | |
| low_cpu_mem_usage=True, | |
| device_map="auto" | |
| ) | |
| else: | |
| self.base_model = AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| torch_dtype=self.dtype, | |
| trust_remote_code=True, | |
| low_cpu_mem_usage=True | |
| ) | |
| self.base_model.to(self.device) | |
| self.base_model.eval() | |
| return self.base_model | |
| def setup_lora(self, trainable_fraction=1.0, r=16, alpha=32, dropout=0.05): | |
| if self.base_model is None: | |
| self.load_base_model() | |
| target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] | |
| peft_config = LoraConfig( | |
| r=r, | |
| lora_alpha=alpha, | |
| target_modules=target_modules, | |
| lora_dropout=dropout, | |
| bias="none", | |
| task_type="CAUSAL_LM" | |
| ) | |
| self.lora_model = get_peft_model(self.base_model, peft_config) | |
| if trainable_fraction < 1.0: | |
| self._apply_trainable_fraction(trainable_fraction) | |
| else: | |
| logger.info("✅ LoRA configured with 100% trainable parameters.") | |
| self.lora_model.print_trainable_parameters() | |
| return self.lora_model | |
| def _apply_trainable_fraction(self, fraction): | |
| lora_params = [(n, p) for n, p in self.lora_model.named_parameters() if "lora_" in n and p.requires_grad] | |
| total_lora = len(lora_params) | |
| num_to_freeze = int(total_lora * (1.0 - fraction)) | |
| if num_to_freeze > 0: | |
| for i, (name, param) in enumerate(lora_params): | |
| if i < num_to_freeze: | |
| param.requires_grad = False | |
| trainable_count = sum(1 for p in self.lora_model.parameters() if p.requires_grad) | |
| logger.info(f"🔒 Froze {num_to_freeze}/{total_lora} LoRA parameters. Remaining trainable: {trainable_count} ({fraction*100}%).") | |
| def autodetect_saved_model(self, output_dir): | |
| output_dir = Path(output_dir) | |
| if output_dir.exists() and (output_dir / "adapter_config.json").exists(): | |
| logger.info(f"✅ Found saved LoRA adapter at {output_dir}. Loading...") | |
| if self.base_model is None: | |
| self.load_base_model() | |
| # NOTE (ZeroGPU compatibility): peft's from_pretrained() infers the load | |
| # device via infer_device(), NOT from the base model's actual .device. | |
| # On HF Spaces ZeroGPU, `spaces` monkey-patches torch.cuda.is_available() | |
| # to report True even before a real GPU is attached to the process, so | |
| # infer_device() always says "cuda" here -- which makes peft try to load | |
| # the adapter's .safetensors weights with a raw, direct CUDA file load. | |
| # That bypasses spaces' safe .to("cuda") staging and crashes with | |
| # "No CUDA GPUs are available". Forcing torch_device="cpu" here loads the | |
| # adapter safely on CPU; the single .to(self.device) call right after is | |
| # a normal nn.Module.to() move, which spaces *does* stage/support. | |
| self.lora_model = PeftModel.from_pretrained(self.base_model, str(output_dir), torch_device="cpu") | |
| self.lora_model.to(self.device) | |
| self.lora_model.eval() | |
| return self.lora_model | |
| logger.info(f"ℹ️ No saved adapter found at {output_dir}.") | |
| return None | |
| def train_lora(self, train_data, output_dir, epochs=3, batch_size=1, grad_accum=4, lr=2e-4, max_length=512): | |
| lr = float(lr) # Ensure lr is a float to prevent PyTorch optimizer TypeError | |
| if self.lora_model is None: | |
| self.setup_lora() | |
| dataset = SpiderDataset(train_data, self.tokenizer, max_length) | |
| def collate_fn(batch): | |
| max_len = max(len(b["input_ids"]) for b in batch) | |
| input_ids, attention_mask, labels = [], [], [] | |
| pad_id = self.tokenizer.pad_token_id | |
| for b in batch: | |
| pad_len = max_len - len(b["input_ids"]) | |
| input_ids.append(b["input_ids"].tolist() + [pad_id] * pad_len) | |
| attention_mask.append(b["attention_mask"].tolist() + [0] * pad_len) | |
| labels.append(b["labels"].tolist() + [-100] * pad_len) | |
| return { | |
| "input_ids": torch.tensor(input_ids, dtype=torch.long).to(self.device), | |
| "attention_mask": torch.tensor(attention_mask, dtype=torch.long).to(self.device), | |
| "labels": torch.tensor(labels, dtype=torch.long).to(self.device) | |
| } | |
| loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn) | |
| optimizer = torch.optim.AdamW( | |
| filter(lambda p: p.requires_grad, self.lora_model.parameters()), | |
| lr=lr | |
| ) | |
| self.lora_model.train() | |
| logger.info(f"🚀 Starting LoRA training | Device: {self.device} | Epochs: {epochs} | Batch: {batch_size} | Grad Accum: {grad_accum} | LR: {lr}") | |
| for epoch in range(epochs): | |
| epoch_loss = 0.0 | |
| progress_bar = tqdm(loader, desc=f"Epoch {epoch+1}/{epochs}", leave=False) | |
| optimizer.zero_grad() | |
| for step, batch in enumerate(progress_bar): | |
| outputs = self.lora_model(**batch) | |
| loss = outputs.loss / grad_accum | |
| loss.backward() | |
| epoch_loss += loss.item() * grad_accum | |
| if (step + 1) % grad_accum == 0: | |
| optimizer.step() | |
| optimizer.zero_grad() | |
| progress_bar.set_postfix({"loss": f"{loss.item() * grad_accum:.4f}"}) | |
| # Clear MPS cache periodically to prevent memory leaks on Mac | |
| if self.device == "mps" and (step + 1) % 10 == 0: | |
| torch.mps.empty_cache() | |
| optimizer.step() | |
| optimizer.zero_grad() | |
| avg_loss = epoch_loss / len(loader) | |
| logger.info(f"✅ Epoch {epoch+1} completed | Avg Loss: {avg_loss:.4f}") | |
| output_dir = Path(output_dir) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| self.lora_model.save_pretrained(str(output_dir)) | |
| self.tokenizer.save_pretrained(str(output_dir)) | |
| logger.info(f"💾 LoRA adapter saved to {output_dir}") | |
| self.lora_model.eval() | |
| return self.lora_model | |
| def load_models(self, lora_path=None): | |
| """Returns a dictionary of models and the tokenizer.""" | |
| if self.tokenizer is None: | |
| self.load_tokenizer() | |
| if self.base_model is None: | |
| self.load_base_model() | |
| models = {"base": self.base_model, "lora": None} | |
| if self.lora_model is not None: | |
| logger.info("Reusing already-loaded LoRA model in memory.") | |
| models["lora"] = self.lora_model | |
| elif lora_path: | |
| lora_path = Path(lora_path) | |
| if lora_path.exists() and (lora_path / "adapter_config.json").exists(): | |
| logger.info(f"Loading LoRA model from {lora_path}") | |
| self.lora_model = PeftModel.from_pretrained(self.base_model, str(lora_path)) | |
| self.lora_model.to(self.device) | |
| self.lora_model.eval() | |
| models["lora"] = self.lora_model | |
| else: | |
| logger.warning(f"LoRA path {lora_path} does not exist or is invalid. Returning base model as 'lora'.") | |
| models["lora"] = self.base_model | |
| return models, self.tokenizer |