model.py

import os
from abc import ABC, abstractmethod

import torch

from huggingface_hub import InferenceClient
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel, PeftConfig
from safetensors.torch import load_file, save_file
from huggingface_hub import hf_hub_download
import tempfile
import shutil
from pathlib import Path


HF_TOKEN = os.environ.get("HF_TOKEN")

# Cache model instances to avoid reloading
_model_cache = {}
model_configs = {
        "baseline": {
                "id": "meta-llama/Llama-3.2-1B-Instruct",
                "type": "endpoint"
        },
        "fine_tuned": {
                "base_model_id": "meta-llama/Llama-3.2-1B-Instruct",
                "adapter_id": "aracape/teaching-assistant-1B-dpo",
                "type": "pipeline" 
        },
        "prompted": {
                "id": "meta-llama/Llama-3.2-1B-Instruct",
                "type": "endpoint"
        },
}

class ModelWrapper(ABC):
    """Abstract base class for model wrappers"""
    
    def __init__(self, model_id: str):
        self.model_id = model_id
    
    @abstractmethod
    def generate(self, messages: list[dict], max_tokens: int, temperature: float) -> str:
        """Generate a response given messages"""
        pass


class InferenceEndpointModel(ModelWrapper):
    """Wrapper for models deployed as HF inference endpoints"""
    
    def __init__(self, model_id: str):
        super().__init__(model_id)
        self.client = InferenceClient(model_id, token=HF_TOKEN)
    
    def generate(self, messages: list[dict], max_tokens: int, temperature: float) -> str:
        """Generate a complete response"""
        result = self.client.chat_completion(
            messages,
            max_tokens=max_tokens,
            temperature=temperature
        )
        return result.choices[0].message["content"]
    
    def generate(self, messages: list[dict], max_tokens: int, temperature: float):
        """Generate a streaming response"""
        response = ""
        for message_chunk in self.client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
        ):
            if message_chunk.choices and message_chunk.choices[0].delta.content:
                token = message_chunk.choices[0].delta.content
                response += token

        return response


class PipelineModel(ModelWrapper):
    """Wrapper for models loaded locally with transformers pipeline"""
    
    def __init__(self, base_model_id: str, adapter_id: str):
        super().__init__(adapter_id)
        # Load pipeline lazily on first use to save memory
        self._pipeline = None
        self.base_model_id = base_model_id
        self.adapter_id = adapter_id

    def _load_model(self):

        if torch.cuda.is_available():
            quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype="bfloat16",
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
            )
        else:
            quantization_config = None

        print(f"Loading base model: {self.base_model_id}")
        base_model = AutoModelForCausalLM.from_pretrained(
            self.base_model_id,
            dtype="auto",
            device_map="auto",
            quantization_config=quantization_config,
            token=HF_TOKEN
        )
        
        # Load the LoRA adapter
        print(f"Loading adapter: {self.adapter_id}")
        model = self._load_adapter_robust(base_model)
        return model

    def _load_adapter_robust(self, base_model):
        """
        Robustly load the adapter, handling potential key mismatches 
        (common with TRL/DPO trained models having extra nesting).
        """
        # Create offload folder if needed for accelerate
        offload_folder = Path("offload")
        offload_folder.mkdir(exist_ok=True)
        
        try:
            return PeftModel.from_pretrained(
                base_model,
                self.adapter_id,
                offload_folder=str(offload_folder),
                token=HF_TOKEN
            )
        except KeyError as e:
            # Check for common key errors indicating extra nesting
            # The error is usually like "KeyError: 'base_model.model.model.model.embed_tokens'"
            # or similar, implying the adapter keys have an extra .model
            if "base_model.model.model" in str(e) or "embed_tokens" in str(e):
                print(f"Encountered KeyError loading adapter: {e}. Attempting to fix keys...")
                return self._load_adapter_with_key_fix(base_model, offload_folder=str(offload_folder))
            raise e

    def _load_adapter_with_key_fix(self, base_model, offload_folder=None):
        """
        Download adapter, fix keys by removing extra 'model' nesting, and load.
        """
        print("Patching adapter weights to fix key mismatch...")
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            try:
                config_path = hf_hub_download(self.adapter_id, "adapter_config.json", token=HF_TOKEN)
                shutil.copy(config_path, temp_path / "adapter_config.json")
            except Exception as e:
                raise RuntimeError(f"Failed to download adapter config: {e}")

            # 2. Download and load weights
            is_safetensors = True
            try:
                weights_path = hf_hub_download(self.adapter_id, "adapter_model.safetensors", token=HF_TOKEN)
                weights = load_file(weights_path)
            except Exception:
                try:
                    weights_path = hf_hub_download(self.adapter_id, "adapter_model.bin", token=HF_TOKEN)
                    weights = torch.load(weights_path, map_location="cpu")
                    is_safetensors = False
                except Exception as e:
                    raise RuntimeError(f"Failed to download adapter weights: {e}")

            # 3. Fix keys: replace 'base_model.model.model.' with 'base_model.model.'
            new_weights = {}
            fixed_count = 0
            for k, v in weights.items():
                if "base_model.model.model." in k:
                    new_k = k.replace("base_model.model.model.", "base_model.model.")
                    new_weights[new_k] = v
                    fixed_count += 1
                else:
                    new_weights[k] = v
            
            print(f"Fixed {fixed_count} keys in adapter weights.")

            # 4. Save fixed weights
            if is_safetensors:
                save_file(new_weights, temp_path / "adapter_model.safetensors")
            else:
                torch.save(new_weights, temp_path / "adapter_model.bin")
            
            # 5. Load PeftModel manually to handle meta tensors
            print("Loading PeftModel with manual state_dict assignment...")
            config = PeftConfig.from_pretrained(temp_dir)
            model = PeftModel(base_model, config)
            
            # Load state dict with assign=True to handle meta tensors
            # This replaces the meta tensors with the loaded CPU tensors
            try:
                model.load_state_dict(new_weights, strict=False, assign=True)
            except TypeError:
                print("Warning: load_state_dict does not support assign=True. Loading might fail for meta tensors.")
                model.load_state_dict(new_weights, strict=False)
                
            return model
    
    @property
    def pipe(self):
        if self._pipeline is None:
            model = self._load_model()
            tokenizer = AutoTokenizer.from_pretrained(self.adapter_id, token=HF_TOKEN)
            self._pipeline = pipeline(
                "text-generation",
                model=model,
                tokenizer=tokenizer,
            )
        return self._pipeline
    
    def generate(self, messages: list[dict], max_tokens: int, temperature: float) -> str:
        """Generate a complete response"""
        result = self.pipe(
            messages,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=temperature > 0,
        )
        return result[0]["generated_text"][-1]["content"]



def create_model(model_key: str) -> ModelWrapper:
    """Factory function to create the appropriate model wrapper"""
    
    config = model_configs[model_key]
    
    if config["type"] == "endpoint":
        return InferenceEndpointModel(config["id"])
    elif config["type"] == "pipeline":
        return PipelineModel(config["base_model_id"], config["adapter_id"])
    else:
        raise ValueError(f"Unknown model type: {config['type']}")


def get_model(model_key: str) -> ModelWrapper:
    """Get or create a model instance"""
    if model_key not in _model_cache:
        _model_cache[model_key] = create_model(model_key)
    return _model_cache[model_key]