models/TimeLanguageModel.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Time Language Model (TLM) for inference.
A multimodal model that combines time series data with language model for time series question answering.
"""
import os
import json
import torch
import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig, AutoTokenizer, AutoModelForCausalLM, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast
from safetensors.torch import load_file
from models.TimeSeriesEncoder import Model
from models.ITFormer import ITFormer
from models.QFormerAdapter import QFormerAdapter
from accelerate import Accelerator

accelerator = Accelerator()

LORA_STATE_MARKERS = (
    ".lora_A.",
    ".lora_B.",
    ".lora_embedding_A.",
    ".lora_embedding_B.",
)


class TLMConfig(PretrainedConfig):
    """Configuration class for Time Language Model."""
    model_type = "vlm_model"
    
    def __init__(self, llm_model_path='LLM/Qwen2.5-0.5B-Instruct',
                 freeze_ts_model=True,
                 ts_pad_num=25,
                 llm_attn_implementation=None,
                 llm_torch_dtype=None,
                 use_lora=False,
                 lora_r=16,
                 lora_alpha=32,
                 lora_dropout=0.05,
                 lora_target_modules=None,
                 gradient_checkpointing=False,
                 **kwargs):
        """Initialize TLM configuration.
        
        Args:
            llm_model_path: Path to the language model
            freeze_ts_model: Whether to freeze time series model parameters
            ts_pad_num: Number of time series padding tokens
            **kwargs: Additional configuration parameters
        """
        self.llm_model_path = llm_model_path
        self.freeze_ts_model = freeze_ts_model
        self.ts_pad_num = ts_pad_num
        self.llm_attn_implementation = llm_attn_implementation
        self.llm_torch_dtype = llm_torch_dtype
        self.use_lora = use_lora
        self.lora_r = lora_r
        self.lora_alpha = lora_alpha
        self.lora_dropout = lora_dropout
        self.lora_target_modules = lora_target_modules or [
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ]
        self.gradient_checkpointing = gradient_checkpointing
        super().__init__(**kwargs)


class TLM(PreTrainedModel, GenerationMixin):
    """Time Language Model for inference."""
    config_class = TLMConfig

    def state_dict(self, *args, **kwargs):
        """Return checkpoint weights without the frozen base LLM.

        The frozen base Qwen weights are reloaded from config.llm_model_path.
        Keep only the trainable LoRA matrices under llm_model.*.
        """
        state_dict = super().state_dict(*args, **kwargs)
        return {
            key: value
            for key, value in state_dict.items()
            if not key.startswith("llm_model.")
            or any(marker in key for marker in LORA_STATE_MARKERS)
        }
    
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, config=None, **kwargs):
        """Load model from pretrained checkpoint.
        
        Args:
            pretrained_model_name_or_path: Path to the checkpoint
            config: Model configuration
            **kwargs: Additional arguments, including ts_config
            
        Returns:
            TLM: Loaded model instance
        """
        if not os.path.exists(pretrained_model_name_or_path):
            raise ValueError(f"Checkpoint path does not exist: {pretrained_model_name_or_path}")

        # Load config.json
        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
        if os.path.exists(config_path):
            with open(config_path, 'r') as f:
                config_dict = json.load(f)
            if config is None:
                config = TLMConfig(**config_dict)
        else:
            if config is None:
                config = TLMConfig()

        # Create model instance with potential ts_config from kwargs
        model = cls(config, **kwargs)

        # Load model weights
        model_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
        if not os.path.exists(model_path):
            model_path = os.path.join(pretrained_model_name_or_path, "model.safetensors")

        state_dict = None
        # 1. Try normal files
        if os.path.exists(model_path):
            if accelerator.is_main_process:
                print(f"Loading model weights from: {model_path}")
            if model_path.endswith('.safetensors'):
                state_dict = load_file(model_path)
            else:
                state_dict = torch.load(model_path, map_location='cpu')
        else:
            # 2. Try split safetensors in the same directory
            all_files = os.listdir(pretrained_model_name_or_path)
            safetensors_files = [f for f in all_files if f.startswith('model-') and f.endswith('.safetensors')]
            safetensors_files.sort()  # Ensure order
            if safetensors_files:
                if accelerator.is_main_process:
                    print(f"Loading split safetensors from: {pretrained_model_name_or_path}")
                state_dict = {}
                for fname in safetensors_files:
                    fpath = os.path.join(pretrained_model_name_or_path, fname)
                    part = load_file(fpath)
                    state_dict.update(part)
                if accelerator.is_main_process:
                    print(f"Successfully loaded {len(safetensors_files)} split safetensors files.")
        if state_dict is not None:
            # Ignore frozen base-LLM weights but retain LoRA matrices.
            ignored_llm_weights = {}
            other_weights = {}
            for k, v in state_dict.items():
                is_lora_weight = any(marker in k for marker in LORA_STATE_MARKERS)
                if k.startswith('llm_model.') and not is_lora_weight:
                    ignored_llm_weights[k] = v
                else:
                    other_weights[k] = v
            if accelerator.is_main_process:
                lora_count = sum(
                    any(marker in key for marker in LORA_STATE_MARKERS)
                    for key in other_weights
                )
                print(f"Found {len(ignored_llm_weights)} frozen LLM weights (will be ignored)")
                print(f"Found {lora_count} LoRA tensors")
                print(f"Found {len(other_weights) - lora_count} non-LLM tensors")
            checkpoint_has_lora = any(
                any(marker in key for marker in LORA_STATE_MARKERS)
                for key in other_weights
            )
            model_has_lora = any(
                "lora_" in name for name, _ in model.llm_model.named_parameters()
            )
            if checkpoint_has_lora and not model_has_lora:
                raise ValueError(
                    "The checkpoint contains LoRA matrices, but the model was "
                    "constructed with use_lora=False."
                )
            if getattr(model.config, "use_lora", False) and not checkpoint_has_lora:
                raise ValueError(
                    "The model was constructed with use_lora=True, but the "
                    "checkpoint does not contain LoRA matrices."
                )
            missing_keys, unexpected_keys = model.load_state_dict(other_weights, strict=False)
            # Filter out LLM-related missing keys since we're not loading LLM weights
            llm_missing_keys = [
                k
                for k in missing_keys
                if k.startswith('llm_model.')
                and not any(marker in k for marker in LORA_STATE_MARKERS)
            ]
            non_llm_missing_keys = [k for k in missing_keys if not k.startswith('llm_model.')]
            missing_lora_keys = [
                k
                for k in missing_keys
                if any(marker in k for marker in LORA_STATE_MARKERS)
            ]
            if llm_missing_keys and accelerator.is_main_process:
                print(f"LLM missing keys (ignored): {len(llm_missing_keys)} keys")
            if missing_lora_keys:
                raise ValueError(f"Missing LoRA checkpoint keys: {missing_lora_keys}")
            if non_llm_missing_keys and accelerator.is_main_process:
                print(f"Non-LLM missing keys: {non_llm_missing_keys}")
            if unexpected_keys and accelerator.is_main_process:
                print(f"Unexpected keys: {unexpected_keys}")
        else:
            if accelerator.is_main_process:
                print(f"Warning: No model weights found at {model_path} or in split safetensors.")

        return model

    def __init__(self, config, ts_config=None):
        """Initialize TLM model.
        
        Args:
            config: TLM configuration
            ts_config: Optional time series configuration (args)
        """
        super().__init__(config)
        self.config = config
        
        if ts_config is None:
            # Create default ts_config if not provided
            class DefaultTSConfig:
                def __init__(self):
                    self.model = 'TimeSeriesEncoder'
                    self.d_model = 512
                    self.n_heads = 8
                    self.e_layers = 4
                    self.patch_len = 60
                    self.stride = 60
                    self.input_len = 600
                    self.dropout = 0.1
                    self.it_d_model = 896
                    self.it_n_heads = 16
                    self.it_layers = 2
                    self.it_dropout = 0.1
                    self.prefix_num = 25
                    self.adapter_type = 'itformer'
            ts_config = DefaultTSConfig()
        
        self.ts_config = ts_config
        
        # 统一属性名对齐逻辑：确保 ts_pad_num 和 prefix_num 存在且一致
        if hasattr(self.ts_config, 'ts_pad_num') and not hasattr(self.ts_config, 'prefix_num'):
            setattr(self.ts_config, 'prefix_num', self.ts_config.ts_pad_num)
        elif hasattr(self.ts_config, 'prefix_num') and not hasattr(self.ts_config, 'ts_pad_num'):
            setattr(self.ts_config, 'ts_pad_num', self.ts_config.prefix_num)
        
        # Initialize LLM model from external path
        try:
            llm_load_kwargs = {}
            attn_impl = getattr(self.config, 'llm_attn_implementation', None)
            dtype_name = getattr(self.config, 'llm_torch_dtype', None)
            dtype_map = {
                "float16": torch.float16,
                "fp16": torch.float16,
                "bfloat16": torch.bfloat16,
                "bf16": torch.bfloat16,
                "float32": torch.float32,
                "fp32": torch.float32,
            }
            if dtype_name:
                normalized_dtype = str(dtype_name).lower()
                if normalized_dtype not in dtype_map:
                    raise ValueError(f"Unsupported llm_torch_dtype: {dtype_name}")
                llm_load_kwargs['torch_dtype'] = dtype_map[normalized_dtype]
            if attn_impl:
                llm_load_kwargs['attn_implementation'] = attn_impl
                # flash_attention_2 / sdpa require fp16/bf16 weights; fp32 errors out.
                if attn_impl in ('flash_attention_2', 'sdpa') and 'torch_dtype' not in llm_load_kwargs:
                    llm_load_kwargs['torch_dtype'] = torch.bfloat16
                if accelerator.is_main_process:
                    print(f"⚡ LLM attention implementation: {attn_impl}")
            llm_load_kwargs['low_cpu_mem_usage'] = True
            self.llm_model = AutoModelForCausalLM.from_pretrained(
                self.config.llm_model_path,
                **llm_load_kwargs,
            )
            self.tokenizer = AutoTokenizer.from_pretrained(self.config.llm_model_path)
            if accelerator.is_main_process:
                print(f"✅ Loaded LLM model from: {self.config.llm_model_path}")
        except Exception as e:
            if accelerator.is_main_process:
                print(f"❌ Failed to load LLM model from {self.config.llm_model_path}: {e}")
            raise e
        
        if self.llm_model is not None:
            self.llm_model.config.pad_token_id = self.tokenizer.pad_token_id

        self._configure_lora()
        
        # Set LLM hidden layer dimension
        ts_config.llm_d_model = self.llm_model.config.hidden_size
        
        # Initialize components
        self.ts_encoder = Model(ts_config)
        
        # 加载预训练的 TS Encoder 权重
        load_path = getattr(ts_config, 'load_ts_encoder', None)
        if load_path and os.path.exists(load_path):
            if accelerator.is_main_process:
                from utils.log_util import adaptive_print
                adaptive_print(f"📥 Loading pre-trained TimeSeries Encoder from: {load_path}")
            
            try:
                if load_path.endswith('.safetensors'):
                    from safetensors.torch import load_file
                    ts_state_dict = load_file(load_path)
                else:
                    ts_state_dict = torch.load(load_path, map_location='cpu')
                
                # 兼容性处理：如果权重包含前缀，进行移除
                new_state_dict = {}
                for k, v in ts_state_dict.items():
                    if k.startswith('model.'):
                        new_state_dict[k[6:]] = v
                    else:
                        new_state_dict[k] = v
                
                msg = self.ts_encoder.load_state_dict(new_state_dict, strict=False)
                if accelerator.is_main_process:
                    adaptive_print(f"✅ TS Encoder weights loaded. Missing: {len(msg.missing_keys)}, Unexpected: {len(msg.unexpected_keys)}")
            except Exception as e:
                if accelerator.is_main_process:
                    adaptive_print(f"❌ Failed to load TS Encoder weights: {e}")
        elif load_path:
            if accelerator.is_main_process:
                from utils.log_util import adaptive_print
                adaptive_print(f"⚠️ Warning: TS Encoder load path '{load_path}' does not exist. Using random initialization.")

        adapter_type = getattr(ts_config, 'adapter_type', 'itformer').lower()
        if adapter_type == 'itformer':
            self.itformer = ITFormer(ts_config)
        elif adapter_type == 'qformer':
            self.itformer = QFormerAdapter(ts_config)
        else:
            raise ValueError(f"Unsupported adapter_type: {adapter_type}")
        if accelerator.is_main_process:
            print(f"🔌 Using adapter: {adapter_type}")
        
        # Projection layers
        self.ts_project = nn.Linear(ts_config.d_model, ts_config.it_d_model)
        self.query_project = nn.Linear(ts_config.llm_d_model, ts_config.it_d_model)
        self.fusion_project = nn.Linear(ts_config.it_d_model, ts_config.llm_d_model)
        
        # 根据配置冻结参数
        self._freeze_layers()

    def _configure_lora(self):
        if not getattr(self.config, "use_lora", False):
            return

        try:
            from peft import LoraConfig, TaskType, get_peft_model
        except ImportError as exc:
            raise RuntimeError("PEFT is required when use_lora=True.") from exc

        target_modules = getattr(self.config, "lora_target_modules", None)
        if isinstance(target_modules, str):
            target_modules = [
                item.strip() for item in target_modules.split(",") if item.strip()
            ]
        if not target_modules:
            raise ValueError("lora_target_modules must not be empty.")

        lora_config = LoraConfig(
            r=int(self.config.lora_r),
            lora_alpha=int(self.config.lora_alpha),
            lora_dropout=float(self.config.lora_dropout),
            bias="none",
            task_type=TaskType.CAUSAL_LM,
            target_modules=list(target_modules),
        )
        self.llm_model = get_peft_model(self.llm_model, lora_config)
        self.llm_model.config.use_cache = False

        if getattr(self.config, "gradient_checkpointing", False):
            self.llm_model.gradient_checkpointing_enable(
                gradient_checkpointing_kwargs={"use_reentrant": False}
            )
            self.llm_model.enable_input_require_grads()

        if accelerator.is_main_process:
            self.llm_model.print_trainable_parameters()

    def _freeze_layers(self):
        """根据配置冻结特定层，保留中间件的可训练性。"""
        # Freeze the base LLM. PEFT has already marked only LoRA matrices as
        # trainable, so preserve those flags when LoRA is enabled.
        if self.llm_model is not None:
            use_lora = bool(getattr(self.config, "use_lora", False))
            for name, param in self.llm_model.named_parameters():
                param.requires_grad = use_lora and "lora_" in name

        # 2. 根据配置冻结 TS Encoder
        if self.config.freeze_ts_model:
            for param in self.ts_encoder.parameters():
                param.requires_grad = False
        else:
            pass

        # 3. 确保中间件是可训练的 (ITFormer 和 Projections)
        # 这些层默认 requires_grad=True，所以不需要额外操作，
        # 除非之前调用了 _setup_inference_mode()

    def _setup_inference_mode(self):
        """Set inference mode, freeze all parameters."""
        for param in self.parameters():
            param.requires_grad = False
        self.eval()
        if accelerator.is_main_process:
            print('🧊 Model set to inference mode - all parameters frozen')

    def eval(self):
        """Set model to evaluation mode."""
        super().eval()
        if self.llm_model is not None:
            self.llm_model.eval()
        if self.ts_encoder is not None:
            self.ts_encoder.eval()
        if self.itformer is not None:
            self.itformer.eval()
        if self.ts_project is not None:
            self.ts_project.eval()
        if self.query_project is not None:
            self.query_project.eval()
        if self.fusion_project is not None:
            self.fusion_project.eval()

    def prepare_inputs_for_generation(self, input_ids, query_ids, past_key_values=None, attention_mask=None, **kwargs):
        """Prepare inputs for text generation.
        
        Args:
            input_ids: Input token IDs
            query_ids: Query token IDs
            past_key_values: Past key values for caching
            attention_mask: Attention mask
            **kwargs: Additional arguments
            
        Returns:
            dict: Prepared inputs for generation
        """
        ts_values = kwargs.get("ts_values", None)
        stage = kwargs.get("stage", None)
        
        if input_ids is None or input_ids.numel() == 0 or ts_values is None or ts_values.numel() == 0:
            return {
                "inputs_embeds": torch.empty(0, self.llm_model.config.hidden_size, device=input_ids.device),
                "attention_mask": attention_mask,
            }
        
        device = next(self.llm_model.parameters()).device
        input_ids = input_ids.to(device)
        ts_values = ts_values.to(device)             
        attention_mask = attention_mask.to(device) 
        
        if ts_values is None:
            raise ValueError("`ts_values` must be provided for generation.")
        
        # Process time series and query
        query_embeds = self.llm_model.get_input_embeddings()(query_ids)
        ts_embeds = self.ts_encoder(ts_values).logits
        ts_embeds = self.ts_project(ts_embeds)
        query_embeds_f = self.query_project(query_embeds)
        it_embeds = self.itformer(query_embeds_f, ts_embeds, stage)
        it_embeds = self.fusion_project(it_embeds)
        
        # Generate inputs_embeds
        inputs_embeds = self.llm_model.get_input_embeddings()(input_ids)
        inputs_embeds = self.merge_input_ids_with_ts_features(it_embeds, inputs_embeds, input_ids)

        return {
            "inputs_embeds": inputs_embeds,
            "attention_mask": attention_mask,
        }

    def forward(self, input_ids=None, query_ids=None, 
                ts_values=None, inputs_embeds=None, stage=None, index=None,
                attention_mask=None, past_key_values=None, labels=None, **kwargs):
        """Forward pass of the model.
        
        Args:
            input_ids: Input token IDs
            query_ids: Query token IDs
            ts_values: Time series values
            inputs_embeds: Pre-computed input embeddings
            stage: Processing stage
            index: Sample index
            attention_mask: Attention mask
            past_key_values: Past key values for caching
            labels: Ground truth labels for loss calculation
            **kwargs: Additional arguments
            
        Returns:
            CausalLMOutputWithPast: Model output
        """
        if inputs_embeds is None:
            # Get query embedding
            query_embeds = self.llm_model.get_input_embeddings()(query_ids)
            # Time series encoding
            ts_embeds = self.ts_encoder(ts_values).logits
            ts_embeds = self.ts_project(ts_embeds)
            query_embeds_f = self.query_project(query_embeds)
            it_embeds = self.itformer(query_embeds_f, ts_embeds, stage)
            it_embeds = self.fusion_project(it_embeds)
            inputs_embeds = self.llm_model.get_input_embeddings()(input_ids)
            inputs_embeds = self.merge_input_ids_with_ts_features(it_embeds, inputs_embeds, input_ids)

        # Forward through LLM
        use_cache = not self.training
        outputs = self.llm_model(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            use_cache=use_cache,
        )
        
        logits = outputs.logits
        return CausalLMOutputWithPast(
            logits=logits,
            past_key_values=outputs.past_key_values if use_cache else None,
        )

    def merge_input_ids_with_ts_features(self, ts_features, inputs_embeds, input_ids):
        batch_size, seq_len, embed_dim = inputs_embeds.shape
        num_tss, num_ts_patches, embed_dim_ = ts_features.shape
        assert embed_dim == embed_dim_, "Embedding dimensions must match."

        pad_token_id = self.tokenizer('<|image_pad|>')['input_ids'][0]
        batch_indices, seq_indices = torch.where(input_ids == pad_token_id)

        if len(batch_indices) != num_tss * num_ts_patches:
            raise ValueError(f"Mismatch: found {len(batch_indices)} pad positions but got {num_tss * num_ts_patches} ts_features.")
        ts_features_flat = ts_features.view(-1, embed_dim).to(
            dtype=inputs_embeds.dtype, 
            device=inputs_embeds.device
        )
        inputs_embeds = inputs_embeds.clone()
        inputs_embeds[batch_indices, seq_indices] = ts_features_flat

        return inputs_embeds