Spaces:
Runtime error
Runtime error
| """ | |
| PEFT Trainer Module | |
| General Parameter-Efficient Fine-Tuning trainer supporting multiple PEFT methods. | |
| """ | |
| from typing import Optional, List, Dict, Any | |
| from pathlib import Path | |
| import torch | |
| from transformers import ( | |
| AutoModelForCausalLM, | |
| AutoTokenizer, | |
| TrainingArguments, | |
| Trainer | |
| ) | |
| from peft import ( | |
| get_peft_model, | |
| PeftConfig, | |
| PeftModel, | |
| prepare_model_for_kbit_training | |
| ) | |
| class PEFTTrainer: | |
| """ | |
| General PEFT Trainer supporting multiple parameter-efficient fine-tuning methods. | |
| Supports: | |
| - LoRA (Low-Rank Adaptation) | |
| - Prefix Tuning | |
| - P-Tuning | |
| - Prompt Tuning | |
| - IA3 (Infused Adapter by Inhibiting and Amplifying Inner Activations) | |
| """ | |
| def __init__( | |
| self, | |
| model_name: str, | |
| peft_config: PeftConfig, | |
| output_dir: str = "./models/peft_output" | |
| ): | |
| """ | |
| Initialize PEFT Trainer. | |
| Args: | |
| model_name: HuggingFace model path or name | |
| peft_config: PEFT configuration (LoraConfig, PrefixTuningConfig, etc.) | |
| output_dir: Directory for saving checkpoints and final model | |
| """ | |
| self.model_name = model_name | |
| self.peft_config = peft_config | |
| self.output_dir = Path(output_dir) | |
| self.output_dir.mkdir(parents=True, exist_ok=True) | |
| self.model = None | |
| self.tokenizer = None | |
| self.trainer = None | |
| def load_model( | |
| self, | |
| use_4bit: bool = False, | |
| use_8bit: bool = False, | |
| device_map: str = "auto" | |
| ) -> None: | |
| """ | |
| Load model with PEFT configuration. | |
| Args: | |
| use_4bit: Use 4-bit quantization | |
| use_8bit: Use 8-bit quantization | |
| device_map: Device mapping strategy | |
| """ | |
| print(f"Loading model: {self.model_name}") | |
| # Load tokenizer | |
| self.tokenizer = AutoTokenizer.from_pretrained( | |
| self.model_name, | |
| trust_remote_code=True | |
| ) | |
| if self.tokenizer.pad_token is None: | |
| self.tokenizer.pad_token = self.tokenizer.eos_token | |
| # Quantization config | |
| quantization_config = None | |
| if use_4bit or use_8bit: | |
| from transformers import BitsAndBytesConfig | |
| if use_4bit: | |
| quantization_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_compute_dtype=torch.float16, | |
| bnb_4bit_use_double_quant=True, | |
| bnb_4bit_quant_type="nf4" | |
| ) | |
| else: | |
| quantization_config = BitsAndBytesConfig(load_in_8bit=True) | |
| # Load base model | |
| self.model = AutoModelForCausalLM.from_pretrained( | |
| self.model_name, | |
| quantization_config=quantization_config, | |
| device_map=device_map, | |
| trust_remote_code=True | |
| ) | |
| # Prepare for k-bit training if quantized | |
| if use_4bit or use_8bit: | |
| self.model = prepare_model_for_kbit_training(self.model) | |
| # Apply PEFT | |
| self.model = get_peft_model(self.model, self.peft_config) | |
| # Print trainable parameters | |
| self.model.print_trainable_parameters() | |
| print("✅ Model loaded with PEFT") | |
| def save_model(self, save_path: Optional[str] = None) -> None: | |
| """ | |
| Save PEFT adapter weights. | |
| Args: | |
| save_path: Path to save adapters | |
| """ | |
| if save_path is None: | |
| save_path = str(self.output_dir / "final_model") | |
| Path(save_path).mkdir(parents=True, exist_ok=True) | |
| self.model.save_pretrained(save_path) | |
| self.tokenizer.save_pretrained(save_path) | |
| print(f"✅ PEFT model saved to: {save_path}") | |
| def load_adapter(self, adapter_path: str) -> None: | |
| """ | |
| Load pre-trained PEFT adapter. | |
| Args: | |
| adapter_path: Path to adapter weights | |
| """ | |
| print(f"Loading PEFT adapter from: {adapter_path}") | |
| self.model = PeftModel.from_pretrained( | |
| self.model, | |
| adapter_path | |
| ) | |
| print("✅ Adapter loaded") | |