# text_encoder.py # Mamba Text Encoder wrapper cho CLIMP-PAR # Dựa trên kiến trúc CLIMP (arXiv:2601.06891): # - Mamba LLM pretrained làm text backbone (thay thế Transformer text encoder) # - Last-token pooling: lấy hidden state ở vị trí token cuối (causal → full context) # - Projection W_t chiếu sang shared embedding space # - Hỗ trợ văn bản dài hơn 77 token (vượt giới hạn CLIP gốc) import os import sys import torch import torch.nn as nn # Thêm đường dẫn mamba local vào sys.path _MAMBA_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'mamba') if _MAMBA_DIR not in sys.path: sys.path.insert(0, _MAMBA_DIR) from mamba.mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel # ============================================================================ # Cấu hình Mamba Text Encoder # ============================================================================ MAMBA_PRETRAINED_MODELS = { # Mamba-1 variants (trên Hugging Face: state-spaces/mamba-*) 'mamba-130m': { 'hf_name': 'state-spaces/mamba-130m-hf', 'd_model': 768, 'n_layer': 24, }, 'mamba-370m': { 'hf_name': 'state-spaces/mamba-370m-hf', 'd_model': 1024, 'n_layer': 48, }, 'mamba-790m': { 'hf_name': 'state-spaces/mamba-790m-hf', 'd_model': 1536, 'n_layer': 48, }, 'mamba-1.4b': { 'hf_name': 'state-spaces/mamba-1.4b-hf', 'd_model': 2048, 'n_layer': 48, }, # Mamba-2 variants 'mamba2-130m': { 'hf_name': 'state-spaces/mamba2-130m', 'd_model': 768, 'n_layer': 24, }, 'mamba2-370m': { 'hf_name': 'state-spaces/mamba2-370m', 'd_model': 1024, 'n_layer': 48, }, 'mamba2-780m': { 'hf_name': 'state-spaces/mamba2-780m', 'd_model': 1536, 'n_layer': 48, }, 'mamba2-1.3b': { 'hf_name': 'state-spaces/mamba2-1.3b', 'd_model': 2048, 'n_layer': 48, }, } class MambaTextEncoder(nn.Module): """ Text Encoder dựa trên Mamba LLM pretrained cho CLIMP-PAR. Theo bài báo CLIMP (Section 3.2 - Text Encoder): - Sử dụng pretrained Mamba LLM (Mamba-1 hoặc Mamba-2) làm text backbone - Last-token pooling: trích hidden state tại vị trí token cuối cùng (non-padding) → Đây là vị trí duy nhất có đầy đủ ngữ cảnh do tính causal của Mamba - Projection W_t chiếu representation sang shared embedding space - Vượt giới hạn 77-token của CLIP nhờ cơ chế autoregressive Công thức (Eq. từ CLIMP): H = [h1, h2, ..., hL] = MambaBackbone(T) text_repr = H[last_non_pad_idx] text_embedding = W_t @ text_repr Args: model_name (str): Tên model Mamba pretrained. Mặc định: 'mamba-130m' embed_dim (int): Chiều shared embedding space. Mặc định: 768 cache_dir (str, optional): Thư mục cache lưu model tải về. freeze_backbone (bool): Đóng băng backbone khi huấn luyện. Mặc định: False """ def __init__(self, model_name='mamba-130m', embed_dim=768, cache_dir=None, freeze_backbone=False): super().__init__() self.model_name = model_name self.embed_dim = embed_dim # Lấy cấu hình if model_name in MAMBA_PRETRAINED_MODELS: config = MAMBA_PRETRAINED_MODELS[model_name] self.d_model = config['d_model'] self.hf_name = config['hf_name'] else: # Fallback: coi model_name là Hugging Face model name trực tiếp self.hf_name = model_name self.d_model = None # Sẽ xác định sau khi load model # Cache directory mặc định if cache_dir is None: cache_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'checkpoints', 'mamba_text') os.makedirs(cache_dir, exist_ok=True) self.cache_dir = cache_dir # 1. Load Mamba backbone + tokenizer từ Hugging Face self.backbone, self.tokenizer = self._load_model_and_tokenizer() # 2. Xác định d_model nếu chưa biết if self.d_model is None: self.d_model = self.backbone.config.d_model # 3. Freeze backbone nếu yêu cầu if freeze_backbone: for param in self.backbone.parameters(): param.requires_grad = False print(f"[MambaText] Backbone đã được đóng băng (freeze).") # 4. Projection head: d_model → embed_dim self.projection = nn.Linear(self.d_model, embed_dim) # 5. Layer norm trước projection (stabilize training) self.norm = nn.LayerNorm(self.d_model) print(f"[MambaText] Khởi tạo thành công: {model_name}") print(f"[MambaText] d_model={self.d_model}, embed_dim={embed_dim}") def _load_model_and_tokenizer(self): """ Tải Mamba model và tokenizer. Xử lý: 1. config.json từ HF chứa key thừa → lọc trước khi tạo MambaConfig 2. HF mới dùng safetensors thay vì pytorch_model.bin → hỗ trợ cả 2 """ try: from transformers import AutoTokenizer from dataclasses import fields as dataclass_fields from huggingface_hub import hf_hub_download print(f"[MambaText] Đang tải model '{self.hf_name}'...") # 1. Tải tokenizer — Mamba-130m dùng GPT-NeoX tokenizer tokenizer = AutoTokenizer.from_pretrained( "EleutherAI/gpt-neox-20b", trust_remote_code=True ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # 2. Tải config.json config_path = hf_hub_download(self.hf_name, "config.json") import json with open(config_path, 'r') as f: config_data = json.load(f) # 3. Lọc bỏ key không thuộc MambaConfig dataclass from mamba_ssm.models.config_mamba import MambaConfig valid_keys = {f.name for f in dataclass_fields(MambaConfig)} filtered_config = {k: v for k, v in config_data.items() if k in valid_keys} removed_keys = set(config_data.keys()) - valid_keys if removed_keys: print(f"[MambaText] Bỏ qua config keys: {removed_keys}") # 4. Tạo model config = MambaConfig(**filtered_config) model = MambaLMHeadModel(config, device="cpu", dtype=torch.float32) # 5. Tải weights — thử safetensors trước, fallback sang pytorch_model.bin try: weights_path = hf_hub_download(self.hf_name, "model.safetensors") from safetensors.torch import load_file state_dict = load_file(weights_path, device="cpu") print(f"[MambaText] Tải weights từ safetensors thành công!") except Exception: try: weights_path = hf_hub_download(self.hf_name, "pytorch_model.bin") state_dict = torch.load(weights_path, map_location="cpu") print(f"[MambaText] Tải weights từ pytorch_model.bin thành công!") except Exception as e2: raise RuntimeError(f"Không tìm thấy weights (safetensors hoặc bin): {e2}") # Đổi tên key cho khớp với local MambaLMHeadModel if "backbone.embeddings.weight" in state_dict: state_dict["backbone.embedding.weight"] = state_dict.pop("backbone.embeddings.weight") # Dùng strict=False vì lm_head.weight được tied với embedding, có thể thiếu trong state_dict model.load_state_dict(state_dict, strict=False) print(f"[MambaText] Tải model thành công!") return model, tokenizer except ImportError as e: raise ImportError( f"Cần cài đặt thư viện: {e}\n" "Chạy: pip install transformers huggingface_hub safetensors" ) except Exception as e: raise RuntimeError(f"Không thể tải model Mamba '{self.hf_name}': {e}") def _get_hidden_states(self, input_ids, attention_mask=None): """ Chạy backbone Mamba và trả về hidden states của layer cuối. Args: input_ids (Tensor): Token IDs, shape (B, L) attention_mask (Tensor, optional): Mask, shape (B, L). 1=valid, 0=padding Returns: hidden_states (Tensor): shape (B, L, d_model) """ # Lấy hidden states từ local mamba_ssm # MambaLMHeadModel trả về CausalLMOutput(logits=lm_logits), ta cần trích xuất hidden states. # Hoặc dùng trực tiếp model.backbone hidden_states = self.backbone.backbone(input_ids, inference_params=None) return hidden_states def _last_token_pooling(self, hidden_states, attention_mask): """ Last-token pooling theo CLIMP. Theo bài báo: "We extract the hidden state at the last non-padding token as the text representation" Lý do: Mamba là mô hình causal (autoregressive), mỗi h_t chỉ chứa thông tin từ t1 đến t_t. Do đó, token cuối cùng (non-padding) là vị trí duy nhất có TOÀN BỘ ngữ cảnh của câu. Args: hidden_states (Tensor): shape (B, L, d_model) attention_mask (Tensor): shape (B, L), 1=valid token, 0=padding Returns: pooled (Tensor): shape (B, d_model) """ if attention_mask is None: # Không có mask → lấy token cuối cùng return hidden_states[:, -1, :] # Tìm vị trí token cuối cùng (non-padding) cho mỗi sample trong batch # attention_mask.sum(1) - 1 = index của token cuối sequence_lengths = attention_mask.sum(dim=1) - 1 # (B,) sequence_lengths = sequence_lengths.long() # Gather hidden state tại vị trí cuối batch_size = hidden_states.shape[0] pooled = hidden_states[torch.arange(batch_size, device=hidden_states.device), sequence_lengths] return pooled def tokenize(self, texts, max_length=None, device=None): """ Tokenize danh sách câu văn bản thành input_ids và attention_mask. Args: texts (list[str]): Danh sách câu prompt max_length (int): Độ dài tối đa để padding/batching. Nếu None sẽ lấy theo câu dài nhất. device: Device đích cho tensor output Returns: dict với 'input_ids' (B, L) và 'attention_mask' (B, L) """ tokens = self.tokenizer( texts, padding='max_length' if max_length else True, truncation=True if max_length else False, max_length=max_length, return_tensors='pt' ) if device is not None: tokens = {k: v.to(device) for k, v in tokens.items()} return tokens def forward(self, input_ids=None, attention_mask=None, texts=None, max_length=None): """ Forward pass: Text → Text Embedding Có thể truyền vào input_ids + attention_mask (đã tokenize sẵn) hoặc truyền texts (list[str]) để tự động tokenize. Args: input_ids (Tensor, optional): Token IDs, shape (B, L) attention_mask (Tensor, optional): Mask, shape (B, L) texts (list[str], optional): Danh sách câu văn bản max_length (int): Độ dài tối đa khi tự tokenize. (Không giới hạn nếu truyền None) Returns: text_features (Tensor): Text embedding, shape (B, embed_dim) """ # Tự động tokenize nếu truyền texts if texts is not None: device = next(self.parameters()).device tokens = self.tokenize(texts, max_length=max_length, device=device) input_ids = tokens['input_ids'] attention_mask = tokens['attention_mask'] assert input_ids is not None, "Phải truyền input_ids hoặc texts" # 1. Lấy hidden states từ Mamba backbone hidden_states = self._get_hidden_states(input_ids, attention_mask) # 2. Last-token pooling → (B, d_model) pooled = self._last_token_pooling(hidden_states, attention_mask) # 3. Layer Norm pooled = self.norm(pooled.float()) # Đảm bảo float32 # 4. Projection → shared embedding space text_features = self.projection(pooled) return text_features def encode_prompts(self, prompt_list, max_length=None, batch_size=32): """ Mã hóa danh sách prompt thành embeddings (hỗ trợ batch lớn). Hữu ích cho việc cache text features của tất cả thuộc tính. Args: prompt_list (list[str]): Danh sách tất cả prompts max_length (int): Độ dài padding batching. None là tự động padding. batch_size (int): Batch size khi encode Returns: all_features (Tensor): shape (num_prompts, embed_dim) """ all_features = [] for i in range(0, len(prompt_list), batch_size): batch_texts = prompt_list[i:i + batch_size] with torch.no_grad(): features = self.forward(texts=batch_texts, max_length=max_length) all_features.append(features) return torch.cat(all_features, dim=0) def get_d_model(self): """Trả về chiều hidden state của backbone Mamba.""" return self.d_model def get_embed_dim(self): """Trả về chiều embedding sau projection.""" return self.embed_dim # ============================================================================ # Factory function # ============================================================================ def create_text_encoder(model_name='mamba-130m', embed_dim=768, **kwargs): """ Factory function tạo Mamba Text Encoder. Args: model_name: Tên model ('mamba-130m', 'mamba-370m', 'mamba2-130m', ...) embed_dim: Chiều embedding đầu ra Returns: MambaTextEncoder instance """ return MambaTextEncoder( model_name=model_name, embed_dim=embed_dim, **kwargs )