Spaces:
Running on Zero
Running on Zero
| from typing import Any, Dict, List, Union | |
| from torch import nn | |
| import torch | |
| from transformers import AutoTokenizer, AutoModel, CLIPTokenizer, CLIPTextModel | |
| import transformers | |
| # Silence the HuggingFace "UNEXPECTED keys" load report and progress bars | |
| transformers.logging.set_verbosity_error() | |
| class TextEncoder(nn.Module): | |
| def __init__(self, config: Dict[str, Any]): | |
| super().__init__() | |
| encoder_type = config["model"]["text_encoder_type"] | |
| pretrained_path = config["model"]["text_encoder_pretrained"][encoder_type] | |
| self.encoder_type = encoder_type | |
| if encoder_type == "clip": | |
| self.tokenizer = CLIPTokenizer.from_pretrained(pretrained_path) | |
| self.model = CLIPTextModel.from_pretrained(pretrained_path) | |
| elif encoder_type in ["biobert", "sbert"]: | |
| self.tokenizer = AutoTokenizer.from_pretrained(pretrained_path) | |
| self.model = AutoModel.from_pretrained(pretrained_path) | |
| else: | |
| raise ValueError(f"Unknown text_encoder_type: {encoder_type}") | |
| for param in self.model.parameters(): | |
| param.requires_grad = False | |
| self.model.eval() | |
| # In-memory embedding cache active ONLY during training mode | |
| self._cache: Dict[str, torch.Tensor] = {} | |
| def clear_cache(self): | |
| self._cache.clear() | |
| def _encode(self, text_list: List[str], device: torch.device) -> torch.Tensor: | |
| inputs = self.tokenizer(text_list, padding=True, truncation=True, return_tensors="pt").to(device) | |
| outputs = self.model(**inputs) | |
| if self.encoder_type == "clip": | |
| # CLIP returns a dedicated pooled_output for the whole sentence | |
| embeddings = outputs.pooler_output | |
| elif self.encoder_type == "biobert": | |
| # BERT models typically use the [CLS] token (index 0) for sentence-level tasks | |
| embeddings = outputs.last_hidden_state[:, 0, :] | |
| elif self.encoder_type == "sbert": | |
| # Sentence-BERT using Mean Pooling across all tokens | |
| attention_mask = inputs['attention_mask'] | |
| token_embeddings = outputs.last_hidden_state | |
| # Expand attention mask to match embedding dimensions | |
| input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() | |
| # Sum the embeddings and divide by the number of non-padded tokens | |
| sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) | |
| sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9) | |
| embeddings = sum_embeddings / sum_mask | |
| return embeddings | |
| def forward(self, text: Union[str, List[str]]) -> torch.Tensor: | |
| device = next(self.model.parameters()).device | |
| is_single = isinstance(text, str) | |
| text_list = [text] if is_single else text | |
| uncached = list({t for t in text_list if t not in self._cache or self._cache[t].device != device}) | |
| if uncached: | |
| new_embeddings = self._encode(uncached, device) | |
| for t, emb in zip(uncached, new_embeddings): | |
| self._cache[t] = emb.detach() | |
| embeddings = torch.stack([self._cache[t] for t in text_list], dim=0) | |
| return embeddings[0] if is_single else embeddings |