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 )