| """Frozen Vietnamese text encoder, standing in for CLIP in T2M-GPT. |
| |
| T2M-GPT conditions its GPT on a single pooled CLIP text embedding. CLIP's text |
| tower is English-only, and this corpus is Vietnamese sign-language gloss, so the |
| tower is swapped for a Vietnamese pretrained encoder. The interface is kept |
| identical to what train_t2m_trans expects: text -> (B, dim) float tensor, so |
| `clip_dim` in the transformer options just becomes this encoder's hidden size. |
| |
| Supported --text-model values are any HF encoder id; two that are relevant here: |
| vinai/phobert-base-v2 (768) Vietnamese RoBERTa, needs word-ish input |
| vinai/bartpho-syllable-base(768) already used by this project's stage-2 baseline |
| """ |
| import torch |
| import torch.nn as nn |
|
|
|
|
| class ViTextEncoder(nn.Module): |
| def __init__(self, model_name="vinai/phobert-base-v2", device="cuda", |
| max_length=64, pooling="mean"): |
| super().__init__() |
| from transformers import AutoModel, AutoTokenizer |
|
|
| self.model_name = model_name |
| self.max_length = max_length |
| self.pooling = pooling |
| self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
| model = AutoModel.from_pretrained(model_name) |
| |
| if hasattr(model, "encoder") and hasattr(model, "decoder"): |
| model = model.encoder |
| self.model = model.to(device).eval() |
| for p in self.model.parameters(): |
| p.requires_grad = False |
| self.device = device |
| self.dim = int(self.model.config.hidden_size) |
|
|
| @torch.no_grad() |
| def forward(self, texts): |
| """list[str] -> (B, dim) float32 on self.device.""" |
| if isinstance(texts, str): |
| texts = [texts] |
| batch = self.tokenizer(list(texts), padding=True, truncation=True, |
| max_length=self.max_length, return_tensors="pt") |
| batch = {k: v.to(self.device) for k, v in batch.items()} |
| out = self.model(**batch).last_hidden_state |
| m = batch["attention_mask"].unsqueeze(-1).float() |
| if self.pooling == "cls": |
| feat = out[:, 0] |
| else: |
| feat = (out * m).sum(1) / m.sum(1).clamp(min=1e-6) |
| return feat.float() |
|
|
| encode = forward |
|
|
