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