src/data/text_encoder.py

import torch
import torch.nn as nn
from transformers import CLIPTextModel, CLIPTokenizer, CLIPConfig
from typing import Optional, Tuple, List
import os


class LightTextEncoder(nn.Module):
    """轻量级文本编码器"""
    def __init__(self, config: dict):
        super().__init__()
        self.config = config
        
        # 编码器参数
        self.vocab_size = config.get('vocab_size', 49408)
        self.hidden_size = config.get('hidden_size', 512)
        self.num_hidden_layers = config.get('num_hidden_layers', 8)
        self.num_attention_heads = config.get('num_attention_heads', 8)
        self.max_position_embeddings = config.get('max_position_embeddings', 77)
        
        # 构建编码器
        self.token_embedding = nn.Embedding(self.vocab_size, self.hidden_size)
        self.position_embedding = nn.Embedding(self.max_position_embeddings, self.hidden_size)
        
        # Transformer层
        self.layers = nn.ModuleList([
            TransformerLayer(self.hidden_size, self.num_attention_heads)
            for _ in range(self.num_hidden_layers)
        ])
        
        self.final_layer_norm = nn.LayerNorm(self.hidden_size)
        
    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        # 嵌入
        token_embeddings = self.token_embedding(input_ids)
        position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0)
        position_embeddings = self.position_embedding(position_ids)
        
        hidden_states = token_embeddings + position_embeddings
        
        # Transformer层
        for layer in self.layers:
            hidden_states = layer(hidden_states, attention_mask)
        
        # 最终层归一化
        hidden_states = self.final_layer_norm(hidden_states)
        
        return hidden_states


class TransformerLayer(nn.Module):
    """Transformer层"""
    def __init__(self, hidden_size: int, num_heads: int):
        super().__init__()
        self.attention = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
        self.attention_norm = nn.LayerNorm(hidden_size)
        
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.GELU(),
            nn.Linear(hidden_size * 4, hidden_size)
        )
        self.mlp_norm = nn.LayerNorm(hidden_size)
        
    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        # 自注意力
        attn_output, _ = self.attention(x, x, x, key_padding_mask=attention_mask)
        x = self.attention_norm(x + attn_output)
        
        # 前馈网络
        mlp_output = self.mlp(x)
        x = self.mlp_norm(x + mlp_output)
        
        return x


class CLIPTextEncoderWrapper:
    """CLIP文本编码器包装器"""
    def __init__(self, model_name: str = 'openai/clip-vit-base-patch32', device: str = 'cuda'):
        self.model_name = model_name
        self.device = device
        
        # 加载tokenizer和模型
        self.tokenizer = CLIPTokenizer.from_pretrained(model_name)
        
        # 只加载文本模型
        self.text_model = CLIPTextModel.from_pretrained(model_name).to(device)
        
        # 冻结参数
        for param in self.text_model.parameters():
            param.requires_grad = False
        
        # 设置为评估模式
        self.text_model.eval()
        
        print(f"已加载CLIP文本编码器: {model_name}")
    
    def encode(self, texts: List[str], return_tensors: str = 'pt') -> torch.Tensor:
        """编码文本"""
        # Tokenize
        inputs = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=77,
            return_tensors=return_tensors
        )
        
        # 移动到设备
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # 编码
        with torch.no_grad():
            outputs = self.text_model(**inputs)
            return outputs.last_hidden_state
    
    def encode_batch(self, texts: List[str], batch_size: int = 32) -> torch.Tensor:
        """分批编码文本"""
        all_embeddings = []
        
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            batch_embeddings = self.encode(batch_texts)
            all_embeddings.append(batch_embeddings.cpu())
        
        return torch.cat(all_embeddings, dim=0)
    
    def save_embeddings(self, texts: List[str], save_path: str):
        """保存文本嵌入"""
        embeddings = self.encode_batch(texts)
        torch.save(embeddings, save_path)
        print(f"嵌入已保存到: {save_path}")


class CachedTextEncoder:
    """带缓存的文本编码器"""
    def __init__(self, encoder, cache_dir: str = './text_cache'):
        self.encoder = encoder
        self.cache_dir = cache_dir
        os.makedirs(cache_dir, exist_ok=True)
        
        # 内存缓存
        self.memory_cache = {}
    
    def encode(self, text: str) -> torch.Tensor:
        """编码文本，使用缓存"""
        # 生成缓存键
        import hashlib
        cache_key = hashlib.md5(text.encode()).hexdigest()
        
        # 检查内存缓存
        if cache_key in self.memory_cache:
            return self.memory_cache[cache_key]
        
        # 检查磁盘缓存
        cache_file = os.path.join(self.cache_dir, f"{cache_key}.pt")
        if os.path.exists(cache_file):
            embedding = torch.load(cache_file)
            self.memory_cache[cache_key] = embedding
            return embedding
        
        # 编码并缓存
        embedding = self.encoder.encode([text])[0]
        
        # 保存到内存缓存
        self.memory_cache[cache_key] = embedding
        
        # 保存到磁盘缓存
        torch.save(embedding, cache_file)
        
        return embedding
    
    def encode_batch(self, texts: List[str]) -> torch.Tensor:
        """批量编码文本"""
        embeddings = []
        
        for text in texts:
            embedding = self.encode(text)
            embeddings.append(embedding.unsqueeze(0))
        
        return torch.cat(embeddings, dim=0)


def create_text_encoder(config: dict) -> CLIPTextEncoderWrapper:
    """创建文本编码器"""
    model_name = config.get('preprocessing', {}).get('tokenizer', 'openai/clip-vit-base-patch32')
    device = config.get('device', 'cuda' if torch.cuda.is_available() else 'cpu')
    
    encoder = CLIPTextEncoderWrapper(model_name, device)
    
    # 如果需要缓存，包装一层
    if config.get('use_cache', True):
        cache_dir = config.get('cache_dir', './data/text_cache')
        encoder = CachedTextEncoder(encoder, cache_dir)
    
    return encoder


def test_text_encoder():
    """测试文本编码器"""
    config = {
        'preprocessing': {
            'tokenizer': 'openai/clip-vit-base-patch32'
        },
        'device': 'cuda' if torch.cuda.is_available() else 'cpu'
    }
    
    encoder = create_text_encoder(config)
    
    # 测试编码
    texts = [
        "A beautiful sunset over the mountains",
        "A cute cat playing with a ball",
        "An astronaut riding a horse on Mars"
    ]
    
    embeddings = encoder.encode(texts)
    
    print(f"文本数量: {len(texts)}")
    print(f"嵌入形状: {embeddings.shape}")
    print(f"嵌入范围: [{embeddings.min():.4f}, {embeddings.max():.4f}]")
    
    return encoder, embeddings


if __name__ == '__main__':
    encoder, embeddings = test_text_encoder()