model_loader.py

"""
Dynamic model loader for any architecture
Supports:
- Classification: config.json + head.safetensors + embeddings.safetensors
- LLM Full: config.json + model.safetensors + tokenizer/
"""
import json
from pathlib import Path
from typing import Dict, Any, Optional, List
import numpy as np
import torch
import torch.nn as nn
from safetensors.torch import load_file as load_safetensors

from config import MODELS_CACHE


# ═══════════════════════════════════════════════════════════════════════════════
# Dynamic Head for Classification
# ═══════════════════════════════════════════════════════════════════════════════

class DynamicHead(nn.Module):
    """
    Dynamic head module that builds architecture from config
    Supports: Linear, Conv2d, BatchNorm, ReLU, Dropout, etc.
    """

    SUPPORTED_LAYERS = {
        "Linear": nn.Linear,
        "Conv1d": nn.Conv1d,
        "Conv2d": nn.Conv2d,
        "BatchNorm1d": nn.BatchNorm1d,
        "BatchNorm2d": nn.BatchNorm2d,
        "LayerNorm": nn.LayerNorm,
        "ReLU": nn.ReLU,
        "GELU": nn.GELU,
        "SiLU": nn.SiLU,
        "Tanh": nn.Tanh,
        "Sigmoid": nn.Sigmoid,
        "Softmax": nn.Softmax,
        "Dropout": nn.Dropout,
        "Dropout2d": nn.Dropout2d,
        "MaxPool1d": nn.MaxPool1d,
        "MaxPool2d": nn.MaxPool2d,
        "AvgPool1d": nn.AvgPool1d,
        "AvgPool2d": nn.AvgPool2d,
        "AdaptiveAvgPool1d": nn.AdaptiveAvgPool1d,
        "AdaptiveAvgPool2d": nn.AdaptiveAvgPool2d,
        "Flatten": nn.Flatten,
        "Embedding": nn.Embedding,
        "LSTM": nn.LSTM,
        "GRU": nn.GRU,
        "MultiheadAttention": nn.MultiheadAttention,
        "TransformerEncoderLayer": nn.TransformerEncoderLayer,
    }

    def __init__(self, config: Dict[str, Any]):
        super().__init__()
        self.config = config
        self.layers = nn.ModuleList()
        self._build_from_config(config)

    def _build_from_config(self, config: Dict[str, Any]):
        """Build network from config"""
        layers_config = config.get("layers", [])

        for layer_cfg in layers_config:
            layer_type = layer_cfg.get("type")
            params = layer_cfg.get("params", {})

            if layer_type not in self.SUPPORTED_LAYERS:
                raise ValueError(f"Unsupported layer type: {layer_type}")

            layer_cls = self.SUPPORTED_LAYERS[layer_type]

            if layer_type in ["ReLU", "GELU", "SiLU", "Tanh", "Sigmoid", "Flatten"]:
                layer = layer_cls(**params) if params else layer_cls()
            else:
                layer = layer_cls(**params)

            self.layers.append(layer)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x)
        return x


# ═══════════════════════════════════════════════════════════════════════════════
# Model Package (Classification)
# ═══════════════════════════════════════════════════════════════════════════════

class ModelPackage:
    """
    Loaded model package with head and embeddings (Classification format)
    """

    def __init__(
        self,
        config: Dict[str, Any],
        head: nn.Module,
        embeddings: Optional[torch.Tensor] = None,
    ):
        self.config = config
        self.head = head
        self.embeddings = embeddings
        self.device = "cpu"
        self.is_llm = False

    def to(self, device: str) -> "ModelPackage":
        self.device = device
        self.head = self.head.to(device)
        if self.embeddings is not None:
            self.embeddings = self.embeddings.to(device)
        return self

    def eval(self) -> "ModelPackage":
        self.head.eval()
        return self

    @torch.no_grad()
    def predict(self, embeddings: torch.Tensor) -> torch.Tensor:
        if embeddings.device != self.device:
            embeddings = embeddings.to(self.device)
        return self.head(embeddings)

    @torch.no_grad()
    def predict_batch(
        self,
        embeddings: np.ndarray,
        batch_size: int = 32,
    ) -> np.ndarray:
        self.head.eval()
        results = []

        for i in range(0, len(embeddings), batch_size):
            batch = torch.tensor(embeddings[i:i+batch_size], dtype=torch.float32)
            batch = batch.to(self.device)
            output = self.head(batch)
            results.append(output.cpu().numpy())

        return np.concatenate(results, axis=0)


# ═══════════════════════════════════════════════════════════════════════════════
# LLM Model Package (Full LLM)
# ═══════════════════════════════════════════════════════════════════════════════

class LLMModelPackage:
    """
    Loaded LLM model package with full model and tokenizer
    """

    def __init__(
        self,
        config: Dict[str, Any],
        model: nn.Module,
        tokenizer: Any,
    ):
        self.config = config
        self.model = model
        self.tokenizer = tokenizer
        self.device = "cpu"
        self.is_llm = True

        # Get dataset info
        self.dataset_name = config.get("dataset", {}).get("name", "unknown")
        self.num_classes = config.get("dataset", {}).get("num_classes", 2)

    def to(self, device: str) -> "LLMModelPackage":
        self.device = device
        self.model = self.model.to(device)
        return self

    def eval(self) -> "LLMModelPackage":
        self.model.eval()
        return self

    @torch.no_grad()
    def generate(self, texts: List[str], max_new_tokens: int = 50) -> List[str]:
        """Generate text completions"""
        self.model.eval()
        results = []

        for text in texts:
            inputs = self.tokenizer(
                text,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=512,
            ).to(self.device)

            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=self.tokenizer.pad_token_id,
            )

            generated = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            results.append(generated)

        return results

    @torch.no_grad()
    def compute_perplexity(self, texts: List[str], batch_size: int = 4) -> float:
        """Compute perplexity on texts"""
        self.model.eval()
        total_loss = 0.0
        total_tokens = 0

        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]

            inputs = self.tokenizer(
                batch_texts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=512,
            ).to(self.device)

            outputs = self.model(**inputs, labels=inputs["input_ids"])

            # Loss is averaged over tokens, multiply by sequence length
            loss = outputs.loss.item()
            num_tokens = inputs["attention_mask"].sum().item()

            total_loss += loss * num_tokens
            total_tokens += num_tokens

        avg_loss = total_loss / total_tokens
        perplexity = torch.exp(torch.tensor(avg_loss)).item()

        return perplexity

    @torch.no_grad()
    def classify_multiple_choice(
        self,
        questions: List[str],
        choices_list: List[List[str]],
        batch_size: int = 4,
    ) -> np.ndarray:
        """
        Classify multiple choice questions.
        Returns index of predicted answer for each question.
        """
        self.model.eval()
        predictions = []

        for question, choices in zip(questions, choices_list):
            choice_losses = []

            for choice in choices:
                prompt = f"{question} {choice}"
                inputs = self.tokenizer(
                    prompt,
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                    max_length=512,
                ).to(self.device)

                outputs = self.model(**inputs, labels=inputs["input_ids"])
                choice_losses.append(outputs.loss.item())

            # Predict the choice with lowest loss
            pred_idx = np.argmin(choice_losses)
            predictions.append(pred_idx)

        return np.array(predictions)

    @torch.no_grad()
    def predict_next_token_probs(
        self,
        texts: List[str],
        target_tokens: List[str],
    ) -> np.ndarray:
        """
        Get probability of target tokens being the next token.
        Used for classification tasks.
        """
        self.model.eval()
        probs = []

        for text, target in zip(texts, target_tokens):
            inputs = self.tokenizer(
                text,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=512,
            ).to(self.device)

            outputs = self.model(**inputs)
            logits = outputs.logits[:, -1, :]  # Last token logits

            # Get probability of target token
            target_ids = self.tokenizer.encode(target, add_special_tokens=False)
            if target_ids:
                target_id = target_ids[0]
                prob = torch.softmax(logits, dim=-1)[0, target_id].item()
            else:
                prob = 0.0

            probs.append(prob)

        return np.array(probs)


# ═══════════════════════════════════════════════════════════════════════════════
# Model Loader
# ═══════════════════════════════════════════════════════════════════════════════

class ModelLoader:
    """
    Loads model packages from local files or IPFS
    Supports both classification (embeddings + head) and LLM (full model) formats
    """

    def __init__(self, cache_dir: Path = MODELS_CACHE):
        self.cache_dir = cache_dir
        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def load_from_directory(self, path: Path) -> ModelPackage | LLMModelPackage:
        """Load model from local directory. Auto-detects format."""
        config_path = path / "config.json"

        if not config_path.exists():
            raise FileNotFoundError(f"config.json not found in {path}")

        with open(config_path, "r") as f:
            config = json.load(f)

        # Check package type
        pkg_type = config.get("type", "classification")

        if pkg_type == "llm_full":
            return self._load_llm_package(path, config)
        else:
            return self._load_classification_package(path, config)

    def _load_classification_package(self, path: Path, config: Dict[str, Any]) -> ModelPackage:
        """Load classification model (embeddings + head)"""
        head_path = path / "head.safetensors"
        embeddings_path = path / "embeddings.safetensors"

        head_config = config.get("head", config)
        head = DynamicHead(head_config)

        if head_path.exists():
            state_dict = load_safetensors(str(head_path))
            self._load_weights_to_head(head, state_dict)

        embeddings = None
        if embeddings_path.exists():
            embeddings_dict = load_safetensors(str(embeddings_path))
            if len(embeddings_dict) == 1:
                embeddings = list(embeddings_dict.values())[0]
            else:
                embeddings = embeddings_dict.get("embeddings")

        return ModelPackage(config, head, embeddings)

    def _load_llm_package(self, path: Path, config: Dict[str, Any]) -> LLMModelPackage:
        """Load full LLM model"""
        model_path = path / "model.safetensors"
        tokenizer_path = path / "tokenizer"

        if not model_path.exists():
            raise FileNotFoundError(f"model.safetensors not found in {path}")

        if not tokenizer_path.exists():
            raise FileNotFoundError(f"tokenizer/ not found in {path}")

        # Get model architecture from config
        model_info = config.get("model", {})
        model_source = model_info.get("source", "gpt2")

        # Load model architecture from HuggingFace
        from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

        # Load config to create empty model
        hf_config = AutoConfig.from_pretrained(model_source, trust_remote_code=True)

        # Create model with config
        model = AutoModelForCausalLM.from_config(hf_config)

        # Load weights from safetensors
        state_dict = load_safetensors(str(model_path))

        # Convert fp16 weights back to model dtype
        for key in state_dict:
            if state_dict[key].dtype == torch.float16:
                state_dict[key] = state_dict[key].float()

        model.load_state_dict(state_dict, strict=False)

        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_path), trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        return LLMModelPackage(config, model, tokenizer)

    def _load_weights_to_head(self, head: DynamicHead, state_dict: Dict[str, torch.Tensor]):
        """Load weights into dynamic head"""
        try:
            head.load_state_dict(state_dict, strict=False)
            return
        except Exception:
            pass

        # Manual mapping for common patterns
        for name, tensor in state_dict.items():
            for i, layer in enumerate(head.layers):
                if hasattr(layer, "weight"):
                    layer_state = layer.state_dict()
                    if "weight" in layer_state:
                        if layer_state["weight"].shape == tensor.shape:
                            layer.weight.data = tensor
                            if "bias" in state_dict.get(name.replace("weight", "bias"), {}):
                                layer.bias.data = state_dict[name.replace("weight", "bias")]
                            break

    def load_from_cache(self, cid: str) -> Optional[ModelPackage | LLMModelPackage]:
        """Load model from cache by CID"""
        cache_path = self.cache_dir / cid
        if cache_path.exists():
            return self.load_from_directory(cache_path)
        return None

    def cache_model(self, cid: str, data: Dict[str, bytes]) -> Path:
        """Cache model data"""
        cache_path = self.cache_dir / cid
        cache_path.mkdir(parents=True, exist_ok=True)

        for filename, content in data.items():
            file_path = cache_path / filename
            with open(file_path, "wb") as f:
                f.write(content)

        return cache_path

    def get_cache_path(self, cid: str) -> Path:
        """Get cache path for CID"""
        return self.cache_dir / cid


# Global loader instance
model_loader = ModelLoader()