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