"""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