oddadmix/Nawah-Router-v3

1"""2Inference for Nawah-Router-v2 — an Arabic zero-shot router with a routing head.3 4The text and every category share one sequence. Each category's character span is mapped to token5indices through the tokenizer's offset mapping and mean-pooled into its own vector; a shared6scorer turns each into one logit, and the softmax runs over the categories actually supplied.7Because the scorer is shared across positions it reads category *content*, not slot index — which8is what makes the label set free text chosen at inference.9 10Layout is text-first, categories-second on purpose. The backbone is causal, so this ordering is11what lets every category token attend to the whole text; reversed, the categories would be12encoded blind to it.13"""14import torch15import torch.nn as nn16from transformers import AutoModel, AutoTokenizer17 18MAX_ROUTES = 919MAX_LENGTH = 32020 21 22class RouterModel(nn.Module):23    def __init__(self, base_model):24        super().__init__()25        self.backbone = AutoModel.from_pretrained(base_model, dtype=torch.float32)26        h = self.backbone.config.hidden_size27        self.score = nn.Sequential(nn.Linear(h, h), nn.GELU(), nn.Linear(h, 1))28        self.config = self.backbone.config29 30    def forward(self, input_ids, attention_mask, cat_pool, n_routes):31        hs = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state32        logits = self.score(torch.bmm(cat_pool.to(hs.dtype), hs)).squeeze(-1)33        ar = torch.arange(logits.size(1), device=logits.device)[None, :]34        return logits.masked_fill(ar >= n_routes[:, None], torch.finfo(logits.dtype).min)35 36    @classmethod37    def from_pretrained(cls, path, token=None):38        import os39        from huggingface_hub import hf_hub_download40        m = cls(path)41        w = (os.path.join(path, "router_model.pt") if os.path.isdir(path)42             else hf_hub_download(path, "router_model.pt", token=token))43        m.load_state_dict(torch.load(w, map_location="cpu", weights_only=True))44        return m.eval()45 46 47def build_text(text, routes):48    head = f"النص:\n{text}\n\nالفئات:\n"49    s, spans = head, []50    for c in routes:51        s += "- "52        spans.append((len(s), len(s) + len(c)))53        s += c + "\n"54    return s, spans55 56 57@torch.no_grad()58def route(model, tok, text, routes):59    """-> [{'route': str, 'score': float}] sorted high to low."""60    routes = [r for r in routes if r and r.strip()][:MAX_ROUTES]61    if not text.strip() or not routes:62        return []63    full, spans = build_text(text, routes)64    enc = tok(full, return_offsets_mapping=True, add_special_tokens=False,65              truncation=True, max_length=MAX_LENGTH)66    ids, offs = enc["input_ids"], enc["offset_mapping"]67    pool = torch.zeros(1, MAX_ROUTES, len(ids))68    for ci, (s, e) in enumerate(spans):69        idx = [t for t, (a, b) in enumerate(offs) if a < e and b > s and a != b]70        if idx:71            pool[0, ci, idx] = 1.0 / len(idx)72    logits = model(torch.tensor([ids]), torch.ones(1, len(ids), dtype=torch.long),73                   pool, torch.tensor([len(routes)]))74    probs = logits.softmax(-1)[0][: len(routes)].tolist()75    out = [{"route": r, "score": p} for r, p in zip(routes, probs)]76    return sorted(out, key=lambda x: -x["score"])77 78 79if __name__ == "__main__":80    M = "oddadmix/Nawah-Router-v2"81    tok = AutoTokenizer.from_pretrained(M)82    model = RouterModel.from_pretrained(M)83    for r in route(model, tok, "الطلب تأخر ساعة والسائق ما رد على الاتصال",84                   ["استفسار عن التوصيل", "شكوى تأخير", "مشكلة في الدفع"]):85        print(f"{r['score']:.3f}  {r['route']}")86