model.py

from dataclasses import dataclass, field
from typing import Optional, Tuple, List, Dict
import json
import os
import shutil

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig, PreTrainedModel, AutoTokenizer
from transformers.modeling_outputs import CausalLMOutputWithPast

try:
    from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
    PEFT_AVAILABLE = True
except ImportError:
    PEFT_AVAILABLE = False


@dataclass
class SiennaConfig:
    """Configuration for Sienna language model.
    
    This is a lightweight wrapper config for training. For HuggingFace compatibility,
    use SiennaHFConfig which extends PretrainedConfig.
    """
    hf_model_name: str = "Qwen/Qwen2.5-0.5B-Instruct"
    # LoRA configuration
    use_lora: bool = False
    lora_r: int = 16
    lora_alpha: int = 32
    lora_dropout: float = 0.05
    lora_target_modules: List[str] = field(default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])
    # Memory optimization
    gradient_checkpointing: bool = False
    use_flash_attention: bool = True


def _qwen_to_sienna_key_mapping(qwen_key: str) -> str:
    """Map Qwen weight keys to Sienna naming convention."""
    # Replace model type identifiers
    key = qwen_key.replace("qwen2", "sienna")
    key = key.replace("Qwen2", "Sienna")
    
    # Rename transformer components to Sienna-specific names
    key = key.replace("model.layers", "sienna_backbone.transformer_blocks")
    key = key.replace("model.embed_tokens", "sienna_backbone.token_embeddings")
    key = key.replace("model.norm", "sienna_backbone.output_norm")
    key = key.replace("lm_head", "sienna_lm_head")
    
    # Rename attention components
    key = key.replace("self_attn", "sienna_attention")
    key = key.replace("q_proj", "query_proj")
    key = key.replace("k_proj", "key_proj")
    key = key.replace("v_proj", "value_proj")
    key = key.replace("o_proj", "output_proj")
    
    # Rename MLP components
    key = key.replace("mlp.gate_proj", "feed_forward.gate_projection")
    key = key.replace("mlp.up_proj", "feed_forward.up_projection")
    key = key.replace("mlp.down_proj", "feed_forward.down_projection")
    
    # Rename normalization layers
    key = key.replace("input_layernorm", "attention_norm")
    key = key.replace("post_attention_layernorm", "ffn_norm")
    
    return key


def _sienna_to_qwen_key_mapping(sienna_key: str) -> str:
    """Reverse mapping: Sienna keys back to Qwen format."""
    key = sienna_key
    
    key = key.replace("sienna_backbone.transformer_blocks", "model.layers")
    key = key.replace("sienna_backbone.token_embeddings", "model.embed_tokens")
    key = key.replace("sienna_backbone.output_norm", "model.norm")
    key = key.replace("sienna_lm_head", "lm_head")
    
    key = key.replace("sienna_attention", "self_attn")
    key = key.replace("query_proj", "q_proj")
    key = key.replace("key_proj", "k_proj")
    key = key.replace("value_proj", "v_proj")
    key = key.replace("output_proj", "o_proj")
    
    key = key.replace("feed_forward.gate_projection", "mlp.gate_proj")
    key = key.replace("feed_forward.up_projection", "mlp.up_proj")
    key = key.replace("feed_forward.down_projection", "mlp.down_proj")
    
    key = key.replace("attention_norm", "input_layernorm")
    key = key.replace("ffn_norm", "post_attention_layernorm")
    
    return key


def _convert_qwen_state_dict_to_sienna(qwen_state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
    """Convert Qwen model weights to Sienna naming convention."""
    sienna_state_dict = {}
    for key, value in qwen_state_dict.items():
        new_key = _qwen_to_sienna_key_mapping(key)
        sienna_state_dict[new_key] = value
    return sienna_state_dict


def save_sienna_tokenizer(hf_model_name: str, output_dir: str = "tokenizer"):
    """Load Qwen tokenizer and save it with Sienna branding.
    
    Args:
        hf_model_name: The HuggingFace model name (e.g., "Qwen/Qwen2.5-0.5B-Instruct")
        output_dir: Directory to save the branded tokenizer (default: "tokenizer")
    """
    print(f"Loading tokenizer from {hf_model_name}...")
    tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
    
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)
    
    # Save tokenizer files first
    tokenizer.save_pretrained(output_dir)
    print(f"✓ Tokenizer files saved to {output_dir}")
    
    # Now rebrand all JSON config files
    json_files = [
        "tokenizer_config.json",
        "special_tokens_map.json",
        "generation_config.json",
    ]
    
    for json_file in json_files:
        json_path = os.path.join(output_dir, json_file)
        if os.path.exists(json_path):
            with open(json_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
            
            # Recursively replace Qwen with Sienna in all string values
            config = _replace_qwen_with_sienna_in_dict(config)
            
            # Add Sienna-specific branding
            if json_file == "tokenizer_config.json":
                config["name_or_path"] = "Sienna-Tokenizer"
                # Remove custom tokenizer class to keep AutoTokenizer happy
                config.pop("tokenizer_class", None)
                # Drop base model reference to avoid Qwen strings in config
                config.pop("_sienna_base_model", None)
            
            with open(json_path, 'w', encoding='utf-8') as f:
                json.dump(config, f, indent=2, ensure_ascii=False)
            
            print(f"  ✓ Rebranded {json_file}")

    # Ensure a minimal config.json exists for AutoTokenizer
    config_path = os.path.join(output_dir, "config.json")
    config_payload = {
        "model_type": "qwen2",
    }
    with open(config_path, 'w', encoding='utf-8') as f:
        json.dump(config_payload, f, indent=2)
    print(f"  ✓ Wrote tokenizer config.json")
    
    print(f"\n✓ Sienna tokenizer successfully created at {output_dir}")
    print(f"  Based on: {hf_model_name}")
    return tokenizer


def _replace_qwen_with_sienna_in_dict(obj):
    """Recursively replace 'Qwen' and 'qwen' with 'Sienna' and 'sienna' in dict/list structures."""
    if isinstance(obj, dict):
        return {key: _replace_qwen_with_sienna_in_dict(value) for key, value in obj.items()}
    elif isinstance(obj, list):
        return [_replace_qwen_with_sienna_in_dict(item) for item in obj]
    elif isinstance(obj, str):
        # Replace Qwen with Sienna (preserve case)
        result = obj.replace("Qwen", "Sienna").replace("qwen", "sienna")
        # Also replace QWen, QWEN variants
        result = result.replace("QWen", "Sienna").replace("QWEN", "SIENNA")
        return result
    else:
        return obj


class Sienna(nn.Module):
    """Sienna Language Model - a custom branded model based on transformer architecture.
    
    This model uses Qwen2.5 architecture internally but presents with Sienna branding
    and renamed layers for HuggingFace deployment.
    """
    def __init__(self, config: SiennaConfig):
        super().__init__()
        self.config = config
        hf_config = AutoConfig.from_pretrained(config.hf_model_name)
        
        # Enable flash attention if requested
        if config.use_flash_attention:
            hf_config.attn_implementation = "flash_attention_2"
        try:
            self.model = AutoModelForCausalLM.from_config(hf_config)
        except (ImportError, RuntimeError) as exc:
            if config.use_flash_attention and _is_flash_attn_error(exc):
                print(f"Flash Attention unavailable, falling back to standard attention: {exc}")
                config.use_flash_attention = False
                hf_config.attn_implementation = "eager"
                self.model = AutoModelForCausalLM.from_config(hf_config)
            else:
                raise
        
        # Enable gradient checkpointing for memory efficiency
        if config.gradient_checkpointing:
            self.model.gradient_checkpointing_enable()
        
        self._is_lora_applied = False
        
        # Store architecture info for HuggingFace
        self.model_type = "sienna"
        self.name_or_path = "Sienna-LM"

    def apply_lora(self) -> None:
        """Apply LoRA adapters for efficient fine-tuning."""
        if not PEFT_AVAILABLE:
            raise ImportError("PEFT is required for LoRA. Install with: pip install peft")
        if self._is_lora_applied:
            return
        
        lora_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            r=self.config.lora_r,
            lora_alpha=self.config.lora_alpha,
            lora_dropout=self.config.lora_dropout,
            target_modules=self.config.lora_target_modules,
            bias="none",
        )
        self.model = get_peft_model(self.model, lora_config)
        self._is_lora_applied = True
        print(f"LoRA applied. Trainable params: {self.model.print_trainable_parameters()}")

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple]]:
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
            past_key_values=past_key_values,
            use_cache=use_cache,
        )
        return outputs.logits, outputs.loss, getattr(outputs, 'past_key_values', None)
    
    def get_trainable_params(self):
        """Return only trainable parameters (useful for LoRA)."""
        return [p for p in self.parameters() if p.requires_grad]
    
    def state_dict(self, *args, **kwargs) -> Dict[str, torch.Tensor]:
        """Get state dict with Sienna-branded layer names."""
        # When LoRA is active we must save the full PEFT state, not just the
        # wrapped base model, otherwise adapter weights are lost on resume.
        base_state = self.model.state_dict(*args, **kwargs)
        return _convert_qwen_state_dict_to_sienna(base_state)
    
    def load_state_dict(self, state_dict: Dict[str, torch.Tensor], strict: bool = True):
        """Load state dict, handling both Sienna and Qwen key formats."""
        if not state_dict:
            return self.model.load_state_dict(state_dict, strict=strict)

        # Check if this is already Sienna format or Qwen format
        sample_key = next(iter(state_dict.keys()))
        
        if "sienna" in sample_key.lower():
            # Convert Sienna keys back to Qwen keys for internal model
            qwen_state = {_sienna_to_qwen_key_mapping(k): v for k, v in state_dict.items()}
        else:
            # Already in Qwen format
            qwen_state = state_dict
        
        # Load into the actual wrapped model. Using only the base model here
        # drops PEFT/LoRA adapter weights during resume.
        return self.model.load_state_dict(qwen_state, strict=strict)
    
    def save_pretrained(self, save_directory: str, save_tokenizer: bool = False, **kwargs):
        """Save model with Sienna branding for HuggingFace.
        
        Args:
            save_directory: Directory to save the model
            save_tokenizer: If True, also saves a Sienna-branded tokenizer to save_directory/tokenizer
        """
        os.makedirs(save_directory, exist_ok=True)

        # Save model weights with Sienna naming
        model_path = os.path.join(save_directory, "pytorch_model.bin")
        torch.save(self.state_dict(), model_path)

        # Create a simple config dict for HuggingFace
        config_dict = {
            "architectures": ["SiennaForCausalLM"],
            "model_type": "sienna",
            "hf_model_name": self.config.hf_model_name,
            "use_lora": self.config.use_lora,
            "lora_r": self.config.lora_r,
            "lora_alpha": self.config.lora_alpha,
            "lora_dropout": self.config.lora_dropout,
            "lora_target_modules": self.config.lora_target_modules,
            "gradient_checkpointing": self.config.gradient_checkpointing,
            "use_flash_attention": self.config.use_flash_attention,
        }
        config_path = os.path.join(save_directory, "config.json")
        with open(config_path, 'w') as f:
            json.dump(config_dict, f, indent=2)

        print(f"✓ Sienna model saved to {save_directory}")
        print(f"  Model architecture: Sienna (based on Qwen2.5)")
        print(f"  Layer naming: Sienna convention for HuggingFace deployment")

        # Optionally save tokenizer
        if save_tokenizer:
            tokenizer_dir = os.path.join(save_directory, "tokenizer")
            save_sienna_tokenizer(self.config.hf_model_name, tokenizer_dir)
            print(f"  ✓ Tokenizer saved to {tokenizer_dir}")

        print(f"\n✓ Sienna model saved successfully")


# Backwards compatibility alias
SiennaModel = Sienna


@torch.no_grad()
def load_sienna_weights(
    model: Sienna, 
    hf_model_name: Optional[str] = None,
    dtype: Optional[torch.dtype] = None,
    model_cache_dir: str = "model",
) -> None:
    """Load Sienna model weights, using cached version if available.
    
    This function first checks if a cached Sienna model exists in the model_cache_dir.
    If found, it loads from cache. Otherwise, it downloads from HuggingFace, 
    transforms to Sienna naming, and saves to cache for future use.
    
    Args:
        model: The Sienna model to load weights into
        hf_model_name: Optional override for the model name (defaults to config's hf_model_name)
        dtype: Optional dtype (torch.bfloat16 recommended for training)
        model_cache_dir: Directory to cache the transformed Sienna model (default: "model")
    """
    load_name = hf_model_name or model.config.hf_model_name
    
    # Check if cached Sienna model exists
    model_path = os.path.join(model_cache_dir, "pytorch_model.bin")
    config_path = os.path.join(model_cache_dir, "config.json")
    tokenizer_path = os.path.join(model_cache_dir, "tokenizer")
    
    if os.path.exists(model_path) and os.path.exists(config_path):
        print(f"✓ Found cached Sienna model in '{model_cache_dir}/'")
        print(f"  Loading pretrained Sienna weights from cache...")
        
        # Load the cached Sienna model
        state_dict = torch.load(model_path, map_location='cpu')
        
        # Apply dtype if specified
        if dtype:
            state_dict = {k: v.to(dtype) if v.is_floating_point() else v for k, v in state_dict.items()}
        
        # Convert sienna keys back to qwen format for internal model
        qwen_state = {_sienna_to_qwen_key_mapping(k): v for k, v in state_dict.items()}
        
        # Load into model
        target_model = model.model.base_model.model if model._is_lora_applied else model.model
        missing, unexpected = target_model.load_state_dict(qwen_state, strict=False)
        
        if missing:
            print(f"  Missing keys: {len(missing)}")
        if unexpected:
            print(f"  Unexpected keys: {len(unexpected)}")
        
        print(f"✓ Successfully loaded cached Sienna model")
        return
    
    # No cache found, download and transform from HuggingFace
    print(f"No cached model found in '{model_cache_dir}/'")
    print(f"Downloading from HuggingFace: {load_name}")
    print(f"This will be cached for future use...")
    
    load_kwargs = {
        "low_cpu_mem_usage": True,
    }
    if dtype:
        load_kwargs["dtype"] = dtype
    
    # Try to use flash attention if available
    if model.config.use_flash_attention:
        load_kwargs["attn_implementation"] = "flash_attention_2"

    try:
        hf_model = AutoModelForCausalLM.from_pretrained(load_name, **load_kwargs)
    except (ImportError, RuntimeError) as exc:
        if model.config.use_flash_attention and _is_flash_attn_error(exc):
            print(f"Flash Attention unavailable while loading weights, retrying with standard attention: {exc}")
            model.config.use_flash_attention = False
            load_kwargs.pop("attn_implementation", None)
            hf_model = AutoModelForCausalLM.from_pretrained(load_name, **load_kwargs)
        else:
            raise
    
    # Get the source state dict
    source_state_dict = hf_model.state_dict()
    
    # Load directly into internal model (no conversion needed since internal model uses Qwen architecture)
    target_model = model.model.base_model.model if model._is_lora_applied else model.model
    missing, unexpected = target_model.load_state_dict(source_state_dict, strict=False)
    
    if missing:
        print(f"Missing keys: {len(missing)}")
    if unexpected:
        print(f"Unexpected keys: {len(unexpected)}")
    
    print(f"✓ Successfully loaded weights from {load_name}")
    
    del hf_model  # Free memory
    
    # Save to cache with Sienna naming
    print(f"\nCaching Sienna model to '{model_cache_dir}/' for future use...")
    os.makedirs(model_cache_dir, exist_ok=True)
    
    # Get state dict with Sienna naming
    sienna_state_dict = model.state_dict()
    
    # Convert back to CPU and original dtype for saving
    if dtype:
        sienna_state_dict = {k: v.cpu().to(torch.float32) if v.is_floating_point() else v.cpu() 
                            for k, v in sienna_state_dict.items()}
    else:
        sienna_state_dict = {k: v.cpu() for k, v in sienna_state_dict.items()}
    
    # Save model weights
    torch.save(sienna_state_dict, model_path)
    print(f"  ✓ Saved model weights: {model_path}")
    
    # Save config
    config_dict = {
        "architectures": ["SiennaForCausalLM"],
        "model_type": "sienna",
        "base_model": load_name,
        "vocab_size": 151936,
        "hidden_size": getattr(model.model.config, 'hidden_size', 896),
        "num_hidden_layers": getattr(model.model.config, 'num_hidden_layers', 24),
        "num_attention_heads": getattr(model.model.config, 'num_attention_heads', 14),
    }
    with open(config_path, 'w') as f:
        json.dump(config_dict, f, indent=2)
    print(f"  ✓ Saved config: {config_path}")
    
    # Save tokenizer with Sienna branding
    save_sienna_tokenizer(load_name, tokenizer_path)
    print(f"  ✓ Saved tokenizer: {tokenizer_path}/")
    
    print(f"\n✓ Sienna model cached successfully!")
    print(f"  Next time you run, it will load instantly from '{model_cache_dir}/'")


# Backwards compatibility alias
def load_qwen_weights(*args, **kwargs):
    """Deprecated: Use load_sienna_weights instead."""
    print("Warning: load_qwen_weights is deprecated, use load_sienna_weights instead")
    return load_sienna_weights(*args, **kwargs)


def _is_flash_attn_error(exc: Exception) -> bool:
    text = str(exc).lower()
    return (
        "flash_attn" in text
        or "flash attention" in text
        or "attn_implementation" in text
        or "undefined symbol" in text
    )