| from __future__ import annotations |
|
|
| import json |
| from collections.abc import Mapping |
|
|
| from torch import Tensor, nn |
|
|
|
|
| def serialize_state(state: Mapping[str, object] | str) -> str: |
| """Serialize structured state in stable key order for the text encoder.""" |
| if isinstance(state, str): |
| return state |
| parts = [] |
| for key in sorted(state): |
| value = state[key] |
| if value is None: |
| continue |
| if isinstance(value, (Mapping, list, tuple)): |
| value = json.dumps(value, ensure_ascii=False, sort_keys=True) |
| parts.append(f"[{key}] {value}") |
| return "\n".join(parts) |
|
|
|
|
| class TransformerTextEncoder(nn.Module): |
| """Shared transformer encoder with attention-mask-aware mean pooling.""" |
|
|
| def __init__( |
| self, model: nn.Module, tokenizer: object, hidden_size: int, max_length: int = 512 |
| ): |
| super().__init__() |
| self.model = model |
| self.tokenizer = tokenizer |
| self.hidden_size = hidden_size |
| self.max_length = max_length |
|
|
| @classmethod |
| def from_pretrained(cls, name: str, max_length: int = 512) -> "TransformerTextEncoder": |
| from transformers import AutoModel, AutoTokenizer |
|
|
| tokenizer = AutoTokenizer.from_pretrained(name) |
| model = AutoModel.from_pretrained(name) |
| return cls(model, tokenizer, model.config.hidden_size, max_length) |
|
|
| def forward(self, texts: list[str]) -> Tensor: |
| batch = self.tokenizer( |
| texts, |
| padding=True, |
| truncation=True, |
| max_length=self.max_length, |
| return_tensors="pt", |
| ) |
| device = next(self.model.parameters()).device |
| batch = {key: value.to(device) for key, value in batch.items()} |
| outputs = self.model(**batch) |
| token_embeddings = outputs.last_hidden_state |
| mask = batch["attention_mask"].unsqueeze(-1).to(token_embeddings.dtype) |
| pooled = (token_embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1) |
| return nn.functional.layer_norm(pooled, (self.hidden_size,)) |
|
|
