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()