File size: 2,290 Bytes
8e5456b | 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 | """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)
# BART-style checkpoints carry a decoder we do not need
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 # (B, L, D)
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
|