milco.py

import torch
import torch.nn as nn
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForMaskedLM,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedModel,
)
from transformers.activations import GELUActivation
from transformers.utils import logging

logger = logging.get_logger(__name__)


class MILCOConfig(PretrainedConfig):
    model_type = "milco"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.lsr_encoder_checkpoint = kwargs.get("lsr_encoder_checkpoint", "bert-base-uncased")
        self.multilingual_encoder_checkpoint = kwargs.get("multilingual_encoder_checkpoint", "bert-base-uncased")
        self.model_variant = kwargs.get("model_variant", "bert")
        # Sub-model dimensions — cached in config.json after first save.
        self.en_vocab_size = kwargs.get("en_vocab_size", None)
        self.en_hidden_size = kwargs.get("en_hidden_size", None)
        self.m_vocab_size = kwargs.get("m_vocab_size", None)
        self.m_hidden_size = kwargs.get("m_hidden_size", None)
        # Sub-model config dicts — stored so __init__ can rebuild architecture offline.
        self.en_sub_config = kwargs.get("en_sub_config", None)
        self.m_sub_config = kwargs.get("m_sub_config", None)

    def _resolve_sub_configs(self):
        """Fetch sub-model configs from HF hub if not already cached."""
        if self.en_sub_config is None:
            en_config = AutoConfig.from_pretrained(self.lsr_encoder_checkpoint, trust_remote_code=True)
            self.en_vocab_size = en_config.vocab_size
            self.en_hidden_size = en_config.hidden_size
            self.en_sub_config = en_config.to_dict()
        if self.m_sub_config is None:
            m_config = AutoConfig.from_pretrained(self.multilingual_encoder_checkpoint, trust_remote_code=True)
            self.m_vocab_size = m_config.vocab_size
            self.m_hidden_size = m_config.hidden_size
            self.m_sub_config = m_config.to_dict()


class MILCOModel(PreTrainedModel):
    config_class = MILCOConfig
    base_model_prefix = "milco"

    def __init__(self, config: MILCOConfig):
        super().__init__(config)
        config._resolve_sub_configs()
        # Build architecture from cached configs — no pretrained weight downloads.
        m_sub_config = AutoConfig.for_model(**config.m_sub_config)
        self.m_model = AutoModel.from_config(m_sub_config)
        en_sub_config = AutoConfig.for_model(**config.en_sub_config)
        en_lsr = AutoModelForMaskedLM.from_config(en_sub_config)
        if hasattr(en_lsr, "cls"):
            self.mlm_head = en_lsr.cls
        else:
            self.mlm_head = en_lsr.lm_head
        self.projector = nn.Linear(config.m_hidden_size, config.en_hidden_size)
        self.activation = GELUActivation()
        self.echo = nn.Linear(config.m_hidden_size, 1)
        self.scale = nn.Parameter(torch.tensor([1.0]))
        self.post_init()

    @property
    def en_tokenizer(self):
        if not hasattr(self, "_en_tokenizer"):
            self._en_tokenizer = AutoTokenizer.from_pretrained(
                self.config.lsr_encoder_checkpoint, trust_remote_code=True,
            )
        return self._en_tokenizer

    @property
    def m_tokenizer(self):
        if not hasattr(self, "_m_tokenizer"):
            self._m_tokenizer = AutoTokenizer.from_pretrained(
                self.config.multilingual_encoder_checkpoint, trust_remote_code=True,
            )
        return self._m_tokenizer

    def get_vocab(self):
        en_term2id = self.en_tokenizer.vocab
        en_id2term = {idx: f"e_{term}" for term, idx in en_term2id.items()}
        m_term2id = self.m_tokenizer.vocab
        m_id2term = {idx + self.config.en_vocab_size: f"m_{term}" for term, idx in m_term2id.items()}
        return en_id2term | m_id2term

    def encode_query(self, *args, **kwargs):
        return self.encode_text(*args, **kwargs)

    def encode_document(self, *args, **kwargs):
        return self.encode_text(*args, **kwargs)

    def encode_text(self, texts, batch_size: int = 32, max_length: int = None, source_view=False, return_dict=False, **kwargs):
        if isinstance(texts, str):
            texts = [texts]
        if max_length is None:
            max_length = self.m_tokenizer.model_max_length
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.to(device)
        reps = []
        for idx in range(0, len(texts), batch_size):
            batch_texts = texts[idx:idx + batch_size]
            inputs = self.m_tokenizer(
                batch_texts, padding=True, truncation=True,
                pad_to_multiple_of=32, max_length=max_length, return_tensors="pt",
            )
            inputs = {k: v.to(device) for k, v in inputs.items()}
            with torch.no_grad(), torch.autocast(device_type=device, dtype=torch.bfloat16):
                output = self._encode(
                    input_ids=inputs["input_ids"],
                    attention_mask=inputs["attention_mask"],
                    source_view=source_view,
                )
            if isinstance(output, tuple):
                sparse_rep = self._build_dual_view_sparse(output)
            else:
                sparse_rep = output.to_sparse()
            reps.append(sparse_rep.cpu())
        combined = sparse_row_cat_2d(reps)
        assert len(texts) == combined.size(0)
        if return_dict:
            return self._sparse_to_dicts(combined)
        return combined

    def _sparse_to_dicts(self, sparse_tensor):
        """Convert a sparse [N, vocab] tensor to a list of {token: weight} dicts."""
        id2term = self.get_vocab()
        sparse_tensor = sparse_tensor.coalesce()
        indices = sparse_tensor.indices()
        values = sparse_tensor.values()
        rows, cols = indices[0], indices[1]
        result = [[] for _ in range(sparse_tensor.size(0))]
        for i in range(len(values)):
            row = rows[i].item()
            col = cols[i].item()
            weight = values[i].item()
            token = id2term.get(col, f"unk_{col}")
            result[row].append((token, weight))
        return [
            dict(sorted(pairs, key=lambda x: x[1], reverse=True))
            for pairs in result
        ]

    def _build_dual_view_sparse(self, output):
        pivot_view, (source_ids, source_weights) = output
        batch_size = source_ids.size(0)
        vocab_size = self.config.en_vocab_size + self.config.m_vocab_size

        pivot_rows, pivot_cols = pivot_view.nonzero(as_tuple=True)
        pivot_vals = pivot_view[pivot_rows, pivot_cols].contiguous()

        source_rows = torch.arange(batch_size, device=source_ids.device).repeat_interleave(source_ids.size(1))
        source_cols = source_ids.view(-1)
        source_vals = source_weights.view(-1)
        nonzero_mask = source_vals > 0
        source_rows = source_rows[nonzero_mask]
        source_cols = source_cols[nonzero_mask] + self.config.en_vocab_size
        source_vals = source_vals[nonzero_mask]

        rows = torch.cat([pivot_rows, source_rows])
        cols = torch.cat([pivot_cols, source_cols])
        indices = torch.stack([rows, cols], dim=0)
        values = torch.cat([pivot_vals, source_vals])
        return torch.sparse_coo_tensor(indices, values, (batch_size, vocab_size))

    def _encode(self, input_ids, attention_mask, source_view=False):
        hidden_states_m = self.m_model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        hidden_states_e = self.activation(self.projector(hidden_states_m))
        logits = self.mlm_head(hidden_states_e) * attention_mask.unsqueeze(-1)
        pivot_view = torch.log1p(torch.relu(logits.max(dim=1).values))
        if source_view:
            token_scores = torch.relu(self.echo(hidden_states_m).squeeze(-1))
            source_weights = token_scores * attention_mask
            return pivot_view, (input_ids, source_weights)
        return pivot_view

    def forward(self, *args, **kwargs):
        return self._encode(*args, **kwargs)

    def save_pretrained(self, save_directory: str, **kwargs):
        kwargs["safe_serialization"] = False
        super().save_pretrained(save_directory, **kwargs)


def sparse_row_cat_2d(tensors):
    """Concatenate a list of 2-D COO sparse tensors along rows (dim=0)."""
    assert len(tensors) > 0, "empty list"
    if len(tensors) == 1:
        return tensors[0]
    tensors = [t.coalesce() for t in tensors]
    assert all(t.layout == torch.sparse_coo for t in tensors), "COO only"

    ncols = tensors[0].size(1)
    assert all(t.dim() == 2 and t.size(1) == ncols for t in tensors), "size(1) mismatch"

    total_rows = sum(t.size(0) for t in tensors)
    device = tensors[0].device

    ind_parts, val_parts = [], []
    row_offset = 0
    for t in tensors:
        idx = t._indices()
        if row_offset:
            idx = idx.clone()
            idx[0] += row_offset
        ind_parts.append(idx)
        val_parts.append(t._values())
        row_offset += t.size(0)

    indices = torch.cat(ind_parts, dim=1)
    values = torch.cat(val_parts, dim=0)
    return torch.sparse_coo_tensor(indices, values, (total_rows, ncols), device=device).coalesce()