| 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)
|
|
|
|
|
| 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
|
|
|
|
|
| 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
|
|
|
|
|
| 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:
|
| """编码文本"""
|
|
|
| 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() |
