modeling_textgeneratorc2k.py

"""Hugging Face `transformers` model and pipeline for TextGeneratorC2K (`storygeneratorc2k` + `Gemma-4-31B` + `EmbeddingGemma-300M` + `userneeds2k` BF16)."""

from __future__ import annotations

import json
import math
import re
import zlib
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
from torch import nn
from transformers import (
    Gemma3TextConfig,
    Gemma3TextModel,
    Gemma4TextConfig,
    Gemma4TextModel,
    Pipeline,
    PretrainedConfig,
    PreTrainedModel,
)


class TextGeneratorC2KConfig(PretrainedConfig):
    model_type = "textgeneratorc2k"

    def __init__(
        self,
        text_config: dict[str, Any] | None = None,
        embedding_config: dict[str, Any] | None = None,
        num_categories: int = 2000,
        num_topics: int = 5,
        num_levels: int = 3,
        subspace_dim: int = 256,
        classifier_input_dim: int = 768,
        classifier_hidden_dim: int = 2048,
        embedding_dense_dim: int = 3072,
        valid_vocab_len: int = 50780,
        runtime_meta_zlib_len: int = 1,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.text_config = text_config or {}
        self.embedding_config = embedding_config or {}
        self.num_categories = num_categories
        self.num_topics = num_topics
        self.num_levels = num_levels
        self.subspace_dim = subspace_dim
        self.classifier_input_dim = classifier_input_dim
        self.classifier_hidden_dim = classifier_hidden_dim
        self.embedding_dense_dim = embedding_dense_dim
        self.valid_vocab_len = valid_vocab_len
        self.runtime_meta_zlib_len = runtime_meta_zlib_len


class MLPClassifier(nn.Module):
    def __init__(self, input_dim: int = 768, hidden_dim: int = 2048, num_labels: int = 2000, dropout: float = 0.15):
        super().__init__()
        self.network = nn.Sequential(
            nn.LayerNorm(input_dim),
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_labels),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)


class TextGeneratorC2KForCausalLM(PreTrainedModel):
    config_class = TextGeneratorC2KConfig
    _tied_weights_keys = []
    _no_split_modules = ["Gemma4DecoderLayer", "Gemma3DecoderLayer", "MLPClassifier"]

    def __init__(self, config: TextGeneratorC2KConfig) -> None:
        super().__init__(config)
        gemma_cfg = Gemma4TextConfig.from_dict(config.text_config)
        eg_cfg = Gemma3TextConfig.from_dict(config.embedding_config)
        self.lm = Gemma4TextModel(gemma_cfg)
        self.embeddinggemma = Gemma3TextModel(eg_cfg)
        self.embedding_dense1 = nn.Linear(config.classifier_input_dim, config.embedding_dense_dim, bias=False)
        self.embedding_dense2 = nn.Linear(config.embedding_dense_dim, config.classifier_input_dim, bias=False)
        self.classifier = MLPClassifier(
            input_dim=config.classifier_input_dim,
            hidden_dim=config.classifier_hidden_dim,
            num_labels=config.num_categories,
        )
        d = gemma_cfg.hidden_size
        c = config.num_categories
        t = config.num_topics
        l = config.num_levels
        s = config.subspace_dim
        self.register_buffer("level_centroids", torch.empty((c, l, d), dtype=torch.bfloat16))
        self.register_buffer("topic_subspace_coeffs", torch.empty((c, t, l, s), dtype=torch.bfloat16))
        self.register_buffer("subspace_basis", torch.empty((s, d), dtype=torch.bfloat16))
        self.register_buffer("sigma_il", torch.empty((c, l), dtype=torch.bfloat16))
        self.register_buffer("embed_mean", torch.empty((d,), dtype=torch.bfloat16))
        self.register_buffer("bridge_proj", torch.empty((d, config.classifier_input_dim), dtype=torch.bfloat16))
        self.register_buffer("valid_vocab_ids", torch.empty((config.valid_vocab_len,), dtype=torch.int32))
        self.register_buffer("runtime_meta_zlib", torch.empty((config.runtime_meta_zlib_len,), dtype=torch.uint8))
        self._runtime_initialized = False
        self.post_init()

    def _ensure_runtime(self) -> None:
        if getattr(self, "_runtime_initialized", False):
            return
        dev = self.lm.embed_tokens.weight.device
        for buf_name in (
            "level_centroids",
            "topic_subspace_coeffs",
            "subspace_basis",
            "sigma_il",
            "embed_mean",
            "bridge_proj",
            "valid_vocab_ids",
        ):
            buf = getattr(self, buf_name)
            if buf.device != dev:
                setattr(self, buf_name, buf.to(dev))
        for mod in (self.embeddinggemma, self.embedding_dense1, self.embedding_dense2, self.classifier):
            if next(mod.parameters()).device != dev:
                mod.to(dev)
        meta_bytes = bytes(self.runtime_meta_zlib.detach().cpu().numpy().tobytes())
        meta = json.loads(zlib.decompress(meta_bytes).decode("utf-8"))
        self.label_ids: list[str] = meta["label_ids"]
        self.label_to_idx = {lid: i for i, lid in enumerate(self.label_ids)}
        self.id_to_token: list[str] = meta["id_to_token"]
        self.token_to_id = {tok: i for i, tok in enumerate(self.id_to_token)}
        self.eos_token_ids = set(meta["eos_token_ids"])
        self.pad_token_id = int(meta["pad_token_id"])
        self.comma_token_id = int(meta["comma_token_id"])
        self.level_lang_prompts = meta["level_lang_prompts"]
        self.cat_leaf_token_ids = meta["cat_leaf_token_ids"]

        cat_file = Path(getattr(self.config, "_name_or_path", ".")) / "catid_define.json"
        if not cat_file.exists():
            cat_file = Path(__file__).resolve().parent / "catid_define.json"
        self.catid_define = json.loads(cat_file.read_text(encoding="utf-8")) if cat_file.exists() else {}

        with torch.inference_mode():
            v_ids = self.valid_vocab_ids.long()
            w_raw = F.normalize(self.lm.embed_tokens.weight[v_ids].float(), dim=-1)
            self.w_centered = F.normalize(w_raw - self.embed_mean.float().unsqueeze(0), dim=-1)
        v_list = v_ids.cpu().tolist()
        self.valid_vocab_set = set(v_list)
        self.valid_words = [self.id_to_token[int(i)].lstrip("▁").lower() for i in v_list]
        self._runtime_initialized = True

    def category_name(self, category_id: str, language: str = "en") -> str:
        self._ensure_runtime()
        return " / ".join(
            self.catid_define.get(c, {}).get(language) or self.catid_define.get(c, {}).get("en", c)
            for c in category_id.split("/")
        )

    def _encode_6lang(self, text: str) -> list[int]:
        s = "▁" + text.strip().replace(" ", "▁")
        ids = [2]  # <bos>
        i, n = 0, len(s)
        while i < n:
            matched = False
            for end in range(min(n, i + 24), i, -1):
                tid = self.token_to_id.get(s[i:end])
                if tid is not None:
                    ids.append(tid)
                    i = end
                    matched = True
                    break
            if not matched:
                i += 1
        return ids[:512]

    @torch.inference_mode()
    def generate_paragraph(
        self,
        category_id: str = "10008/20081/31194",
        language: str = "en",
        level: int = 1,
        topic_idx: int = 1,
        gamma: float = 0.45,
        seed: int = 42,
        max_new_tokens: int = 260,
    ) -> dict[str, Any]:
        self._ensure_runtime()
        dev = self.lm.embed_tokens.weight.device
        cat_idx = self.label_to_idx[category_id]
        lvl_idx = max(0, min(2, level - 1))
        t_idx = max(0, min(4, topic_idx - 1))

        ebar = self.level_centroids[cat_idx, lvl_idx].float()
        coeff = self.topic_subspace_coeffs[cat_idx, t_idx, lvl_idx].float()
        basis = self.subspace_basis.float()
        e_anchor = F.normalize(ebar + torch.matmul(coeff, basis), dim=-1)

        g = torch.Generator(device=dev).manual_seed(seed + cat_idx * 31 + level * 7 + topic_idx)
        eps = torch.randn((basis.shape[0],), generator=g, device=dev, dtype=torch.float32) / math.sqrt(basis.shape[0])
        delta = self.sigma_il[cat_idx, lvl_idx].float() * torch.matmul(eps, basis)
        h_tilde = F.normalize(e_anchor + gamma * delta, dim=-1)

        sims = torch.matmul(self.w_centered, h_tilde)
        _, top_pos = torch.topk(sims, k=64)
        c_ids: list[int] = []
        c_words: list[str] = []
        seen: set[str] = set()
        for pos in top_pos.cpu().tolist():
            w = self.valid_words[pos]
            if len(w) < 4 or any(w[:4] == s[:4] for s in seen):
                continue
            seen.add(w)
            c_ids.append(int(self.valid_vocab_ids[pos].item()))
            c_words.append(w)
            if len(c_ids) >= 7:
                break

        tmpl = self.level_lang_prompts[language][str(level)]
        cat_toks = self.cat_leaf_token_ids[language][cat_idx]
        seq = list(tmpl["prefix_ids"]) + list(cat_toks) + list(tmpl["mid_ids"])
        for i, cid in enumerate(c_ids):
            if i > 0:
                seq.append(self.comma_token_id)
            seq.append(cid)
        seq.extend(tmpl["suffix_ids"])

        inp = torch.tensor([seq], device=dev, dtype=torch.long)
        mask = torch.ones_like(inp)
        out = self.lm(inputs_embeds=self.lm.embed_tokens(inp), attention_mask=mask, use_cache=True)
        pkv, last_h = out.past_key_values, out.last_hidden_state[:, -1, :]
        gen: list[int] = []
        for _ in range(max_new_tokens):
            logits = F.linear(last_h.to(dev), self.lm.embed_tokens.weight).float()
            cap = getattr(self.lm.config, "final_logit_softcapping", None)
            if cap:
                logits = torch.tanh(logits / cap) * cap
            nid = int(torch.argmax(logits, dim=-1).item())
            if nid in self.eos_token_ids:
                break
            gen.append(nid)
            nxt = torch.tensor([[nid]], device=dev, dtype=torch.long)
            mask = torch.cat([mask, torch.ones((1, 1), device=dev, dtype=torch.long)], dim=1)
            out = self.lm(inputs_embeds=self.lm.embed_tokens(nxt), attention_mask=mask, past_key_values=pkv, use_cache=True)
            pkv, last_h = out.past_key_values, out.last_hidden_state[:, -1, :]

        text = re.sub(r"\s+", " ", "".join(self.id_to_token[t] for t in gen).replace("▁", " ")).strip()
        if language == "cn":
            text = re.sub(r"\s*,\s*", "，", text)
            text = re.sub(r"\s*;\s*", "；", text)
            text = re.sub(r"\s*:\s*", "：", text)
            text = re.sub(r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])", "", text)
        elif language == "jp":
            text = re.sub(r"\s*,\s*", "、", text)
            text = re.sub(r"(?<=[\u3040-\u30ff\u4e00-\u9fff])\s+(?=[\u3040-\u30ff\u4e00-\u9fff])", "", text)

        leaf_name = self.catid_define.get(category_id.split("/")[-1], {}).get(language) or self.category_name(category_id, "en")
        return {
            "category_id": category_id,
            "category_name": self.category_name(category_id, language),
            "language": language,
            "level": level,
            "concepts": c_words,
            "paragraph": text,
            "categorical_scores": self.classify_text(text, query=leaf_name, language=language, top_k=5)["categorical_scores"],
        }

    @torch.inference_mode()
    def classify_text(self, text: str, query: str = "", language: str = "en", top_k: int = 5) -> dict[str, Any]:
        self._ensure_runtime()
        dev = self.lm.embed_tokens.weight.device
        doc_str = f"Query: {query}\nSnippet: {text}" if query else f"Query: {text[:80]}\nSnippet: {text}"
        inp = torch.tensor([self._encode_6lang(doc_str)], device=dev, dtype=torch.long)
        eg_out = self.embeddinggemma(input_ids=inp)
        pooled = eg_out.last_hidden_state.mean(dim=1)
        dense = self.embedding_dense2(self.embedding_dense1(pooled))
        emb_768 = F.normalize(dense.float(), dim=-1).to(self.embedding_dense1.weight.dtype)
        eg_probs = self.classifier(emb_768).detach().float().sigmoid()[0]

        words = re.findall(r"[A-Za-z]{4,}", text.lower())
        c_ids: list[int] = []
        seen_ids: set[int] = set()
        for w in words:
            tid = self.token_to_id.get(f"▁{w}")
            if tid is not None and tid in self.valid_vocab_set and tid not in seen_ids:
                seen_ids.add(tid)
                c_ids.append(tid)
        if not c_ids:
            for tid in inp[0].cpu().tolist()[1:]:
                if tid in self.valid_vocab_set and tid not in seen_ids:
                    seen_ids.add(tid)
                    c_ids.append(tid)
        if c_ids:
            t_ids = torch.tensor(c_ids[:256], device=dev, dtype=torch.long)
            w_tok = F.normalize(self.lm.embed_tokens.weight[t_ids].float(), dim=-1)
            w_c = F.normalize(w_tok - self.embed_mean.float().unsqueeze(0), dim=-1)
            tok_cat_sims = torch.einsum("cld,md->clm", self.level_centroids.float(), w_c)
            k_top = min(6, w_c.shape[0])
            bank_sims = torch.topk(tok_cat_sims, k=k_top, dim=-1).values.mean(dim=-1).max(dim=1).values.clamp_min(0.0)
            scores = 0.65 * eg_probs + 0.35 * bank_sims
        else:
            scores = eg_probs
        scores[0] = 0.0
        vals, idxs = torch.topk(scores, k=min(top_k, len(scores)))
        preds = [
            {
                "category_id": self.label_ids[idx],
                "category_name": self.category_name(self.label_ids[idx], language),
                "score": round(float(val), 4),
                "userneeds2k_prob": round(float(eg_probs[idx].item()), 4),
            }
            for val, idx in zip(vals.cpu().tolist(), idxs.cpu().tolist())
        ]
        return {
            "text": text,
            "categorical_scores": preds,
            "predictions": preds,
        }


class TextGeneratorC2KPipeline(Pipeline):
    _load_tokenizer = False
    _load_image_processor = False
    _load_video_processor = False
    _load_feature_extractor = False
    _load_processor = False

    def _sanitize_parameters(self, **kwargs: Any) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
        return {}, kwargs, {}

    def preprocess(self, inputs: Any) -> Any:
        return inputs

    def _forward(self, model_inputs: Any, **forward_params: Any) -> dict[str, Any]:
        if isinstance(model_inputs, dict) and "category_id" in model_inputs:
            merged = {**forward_params, **model_inputs}
            return self.model.generate_paragraph(**merged)
        if isinstance(model_inputs, str):
            self.model._ensure_runtime()
            if model_inputs in self.model.label_to_idx:
                return self.model.generate_paragraph(category_id=model_inputs, **forward_params)
            return self.model.classify_text(model_inputs, **forward_params)
        raise ValueError(f"Unsupported pipeline input: {model_inputs}")

    def postprocess(self, model_outputs: dict[str, Any]) -> dict[str, Any]:
        return model_outputs


# Aliases for StoryGeneratorC2K
StoryGeneratorC2KConfig = TextGeneratorC2KConfig
StoryGeneratorC2KForCausalLM = TextGeneratorC2KForCausalLM
StoryGeneratorC2KPipeline = TextGeneratorC2KPipeline