Update

app.py

@@ -4,6 +4,7 @@ import pickle
 from huggingface_hub import hf_hub_download
 from Nested.nn.BertSeqTagger import BertSeqTagger
 from transformers import AutoTokenizer, AutoModel
+import inspect
 app = FastAPI()
 print("Version 2...")
 
@@ -14,10 +15,9 @@ print("Version 2...")
 # )
 
 
-pretrained_path = "aubmindlab/bert-base-arabertv2"  # set to your training backbone
+pretrained_path = "aubmindlab/bert-base-arabertv2"  # must match training
 tokenizer = AutoTokenizer.from_pretrained(pretrained_path)
-encoder = AutoModel.from_pretrained(pretrained_path)
-encoder.eval()
+encoder = AutoModel.from_pretrained(pretrained_path).eval()
 
 checkpoint_path = hf_hub_download(
     repo_id="SinaLab/Nested",
@@ -80,7 +80,7 @@ ckpt = torch.load(checkpoint_path, map_location="cpu")
 model = load_model_from_checkpoint(model, ckpt, strict=False)
 # model.eval()
 
-def predict_ner_with_external_encoder(sentence, tagger, encoder, tokenizer, id2label, device="cpu", max_length=128):
+def predict_ner(sentence, tagger, encoder, tokenizer, id2label, device="cpu", max_length=128):
     tagger.to(device).eval()
     encoder.to(device).eval()
 
@@ -101,21 +101,11 @@ def predict_ner_with_external_encoder(sentence, tagger, encoder, tokenizer, id2l
             attention_mask=enc.get("attention_mask", None)
         ).last_hidden_state  # [1, seq_len, hidden]
 
-        ignore_idx = getattr(tagger, "label_ignore_idx", 0)
-        dummy_labels = torch.full((1, x.size(1)), ignore_idx, dtype=torch.long, device=device)
-
-        out = tagger(x, dummy_labels)
-
-        if isinstance(out, (tuple, list)):
-            logits = out[-1]
-        elif hasattr(out, "logits"):
-            logits = out.logits
-        else:
-            logits = out
+        logits = _call_tagger(tagger, x, device)
 
     pred_ids = logits.argmax(dim=-1)[0].tolist()
-
-    word_ids = tokenizer(words, is_split_into_words=True, truncation=True, max_length=max_length).word_ids()
+    word_ids = _get_word_ids(tokenizer, words, tokenizer(words, is_split_into_words=True, return_tensors="pt",
+                                                        truncation=True, max_length=max_length), max_length)
 
     results = []
     seen = set()
@@ -127,7 +117,6 @@ def predict_ner_with_external_encoder(sentence, tagger, encoder, tokenizer, id2l
 
     return results
 
-
 def find_label_vocab(vocabs):
     for i, v in enumerate(vocabs):
         if hasattr(v, "itos"):
@@ -137,6 +126,37 @@ def find_label_vocab(vocabs):
     return None, None
 
 
+def _get_word_ids(tokenizer, words, enc, max_length):
+    # Fast tokenizers: BatchEncoding has word_ids()
+    if hasattr(enc, "word_ids"):
+        return enc.word_ids(batch_index=0)
+    # Fallback
+    return tokenizer(words, is_split_into_words=True, truncation=True, max_length=max_length).word_ids()
+
+def _call_tagger(tagger, x, device):
+    # Calls forward in a compatible way (x only vs x+labels, etc.)
+    params = list(inspect.signature(tagger.forward).parameters.keys())
+    # common: ['x']  or  ['x','labels',...]
+    if "labels" in params:
+        ignore_idx = getattr(tagger, "label_ignore_idx", 0)
+        labels = torch.full((x.size(0), x.size(1)), ignore_idx, dtype=torch.long, device=device)
+
+        kwargs = {}
+        if "segments_mask" in params:
+            kwargs["segments_mask"] = None
+        if "get_sent_repr" in params:
+            kwargs["get_sent_repr"] = False
+
+        out = tagger(x, labels, **kwargs)
+    else:
+        out = tagger(x)
+
+    # normalize outputs to logits tensor
+    if isinstance(out, (tuple, list)):
+        return out[-1]
+    if hasattr(out, "logits"):
+        return out.logits
+    return out
 
 
 label_vocab = label_vocab[0]  # the list loaded from pickle
@@ -150,5 +170,5 @@ id2label = {i: s for i, s in enumerate(label_vocab.itos)}
 sentence = "ذهب احمد الى السوق"
 # id2label = {i: s for i, s in enumerate(label_vocab.itos)}
 # pairs = predict_ner(sentence, model, label_vocab, device="cpu")
-pairs = predict_ner_with_external_encoder(sentence, model, encoder, tokenizer, id2label, device="cpu")
+pairs = predict_ner(sentence, model, encoder, tokenizer, id2label, device="cpu")
 print(pairs)

requirements.txt

@@ -3,4 +3,5 @@ fastapi
 uvicorn
 numpy
 huggingface_hub
-transformers
+transformers
+inspect