Spaces:
Running on Zero
Running on Zero
File size: 3,406 Bytes
7fe2a79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | 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
@torch.inference_mode()
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 |