| """ |
| Qwen Text Encoder |
| Qwen 文本编码器 - 使用 Qwen3 模型进行文本编码 |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| from typing import List, Optional, Union, Tuple |
|
|
|
|
| def load_qwen_model(model_path: str, device: str = "cuda"): |
| """ |
| Load Qwen3 embedding model |
| 加载 Qwen3 嵌入模型 |
| """ |
| try: |
| from sentence_transformers import SentenceTransformer |
| model = SentenceTransformer(model_path) |
| model.to(device) |
| return model |
| except ImportError: |
| print("Warning: sentence-transformers not available. Using mock embeddings.") |
| return None |
|
|
|
|
| def encode_text_with_qwen( |
| qwen_model, |
| texts: List[str], |
| device: str = "cuda", |
| max_length: int = 512, |
| use_query_mode: bool = False |
| ) -> torch.Tensor: |
| """ |
| Encode text using Qwen3 model |
| 使用 Qwen3 模型编码文本 |
| Args: |
| qwen_model: Qwen3 embedding model |
| texts: List of text strings to encode |
| device: Device to run on |
| max_length: Maximum sequence length |
| use_query_mode: Whether to use query prompt for better understanding |
| """ |
| if qwen_model is None: |
| |
| batch_size = len(texts) |
| return torch.randn(batch_size, 1024, device=device, dtype=torch.float32) |
| |
| with torch.no_grad(): |
| |
| embeddings = qwen_model.encode( |
| texts, |
| prompt_name="query" if use_query_mode else None, |
| convert_to_tensor=True, |
| device=device, |
| max_seq_length=max_length, |
| output_value="token_embeddings" if not use_query_mode else "sentence_embedding" |
| ) |
| |
| return embeddings if use_query_mode else torch.stack(embeddings, dim=0) |
| |
|
|
| class QwenTextEncoder(nn.Module): |
| """ |
| Qwen Text Encoder wrapper for training and inference |
| 用于训练和推理的 Qwen 文本编码器包装器 |
| """ |
| |
| def __init__( |
| self, |
| model_path: str = "models/Qwen3-Embedding-0.6B", |
| device: str = "cuda", |
| max_length: int = 512, |
| freeze_encoder: bool = True |
| ): |
| super().__init__() |
| self.device = device |
| self.max_length = max_length |
| self.freeze_encoder = freeze_encoder |
| |
| |
| self.qwen_model = load_qwen_model(model_path, device) |
| |
| |
| if self.freeze_encoder and self.qwen_model is not None: |
| for param in self.qwen_model.parameters(): |
| param.requires_grad = False |
| |
| def encode_prompts( |
| self, |
| prompts: List[str], |
| negative_prompts: Optional[List[str]] = None, |
| do_classifier_free_guidance: bool = False |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| """ |
| Encode prompts using Qwen3 model |
| 使用 Qwen3 模型编码提示词 |
| |
| Returns: |
| tuple: (text_embeddings, pooled_embeddings) |
| - text_embeddings: [batch_size, 1024] for sequence embeddings |
| - pooled_embeddings: [batch_size, 1024] for pooled embeddings |
| """ |
| batch_size = len(prompts) |
| |
| |
| text_embeddings = encode_text_with_qwen( |
| self.qwen_model, prompts, self.device, |
| max_length=self.max_length, use_query_mode=False |
| ) |
| |
| |
| pooled_embeddings = encode_text_with_qwen( |
| self.qwen_model, prompts, self.device, |
| max_length=self.max_length, use_query_mode=True |
| ) |
| |
| |
| if do_classifier_free_guidance: |
| if negative_prompts is None: |
| negative_prompts = [""] * batch_size |
| |
| |
| negative_text_embeddings = encode_text_with_qwen( |
| self.qwen_model, negative_prompts, self.device, |
| max_length=self.max_length, use_query_mode=False |
| ) |
| |
| negative_pooled_embeddings = encode_text_with_qwen( |
| self.qwen_model, negative_prompts, self.device, |
| max_length=self.max_length, use_query_mode=True |
| ) |
| |
| |
| text_embeddings = torch.cat([negative_text_embeddings, text_embeddings], dim=0) |
| pooled_embeddings = torch.cat([negative_pooled_embeddings, pooled_embeddings], dim=0) |
| |
| return text_embeddings, pooled_embeddings |
| |
| def forward(self, prompts: List[str], negative_prompts: Optional[List[str]] = None): |
| """ |
| Forward pass for text encoding |
| Args: |
| prompts: List of text prompts |
| negative_prompts: Optional list of negative prompts |
| Returns: |
| tuple: (text_embeddings, pooled_embeddings) |
| """ |
| return self.encode_prompts(prompts, negative_prompts, do_classifier_free_guidance=(negative_prompts is not None)) |
| |
| def train(self, mode: bool = True): |
| """Override train mode to handle frozen encoder""" |
| super().train(mode) |
| if self.freeze_encoder and self.qwen_model is not None: |
| self.qwen_model.eval() |
| return self |
|
|
