File size: 2,068 Bytes
3155e2a | 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 | 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,))
|