Finalll

app.py

@@ -3,7 +3,7 @@ import torch
 import pickle
 from huggingface_hub import hf_hub_download
 from Nested.nn.BertSeqTagger import BertSeqTagger
-
+from transformers import AutoTokenizer
 app = FastAPI()
 print("Version 2...")
 
@@ -13,6 +13,9 @@ print("Version 2...")
     # filename="tag_vocab.pkl"
 # )
 
+pretrained_path = "aubmindlab/bert-base-arabertv2"  # change if different in your training
+tokenizer = AutoTokenizer.from_pretrained(pretrained_path)
+
 checkpoint_path = hf_hub_download(
     repo_id="SinaLab/Nested",
     filename="checkpoints/checkpoint_2.pt"
@@ -74,22 +77,18 @@ ckpt = torch.load(checkpoint_path, map_location="cpu")
 model = load_model_from_checkpoint(model, ckpt, strict=False)
 # model.eval()
 
-def predict_ner(sentence: str, model, id2label: dict, device="cpu"):
+def predict_ner(sentence: str, model, tokenizer, id2label: dict, device="cpu", max_length=128):
     model.to(device)
     model.eval()
 
     words = sentence.split()
 
-    tokenizer = getattr(model, "tokenizer", None)
-    if tokenizer is None:
-        raise ValueError("Model has no tokenizer. Use AutoTokenizer and attach it or pass it explicitly.")
-
     enc = tokenizer(
         words,
         is_split_into_words=True,
         return_tensors="pt",
         truncation=True,
-        padding=False
+        max_length=max_length
     )
     enc = {k: v.to(device) for k, v in enc.items()}
 
@@ -98,28 +97,27 @@ def predict_ner(sentence: str, model, id2label: dict, device="cpu"):
             out = model(**enc)
             logits = out.logits if hasattr(out, "logits") else out
         except TypeError:
-            if not hasattr(model, "transformer") or not hasattr(model, "classification_head"):
-                raise
+            # fallback for your custom BertSeqTagger-like model
             h = model.transformer(**enc).last_hidden_state
-            h = model.dropout(h) if hasattr(model, "dropout") else h
+            if hasattr(model, "dropout"):
+                h = model.dropout(h)
             logits = model.classification_head(h)
 
-    pred_ids = logits.argmax(dim=-1).squeeze(0).tolist()
+    pred_ids = logits.argmax(dim=-1)[0].tolist()
 
-    word_ids = enc.get("input_ids").new_tensor([0])  # placeholder to keep structure
-    word_ids = tokenizer(words, is_split_into_words=True).word_ids()
+    word_ids = enc["input_ids"].new_zeros(enc["input_ids"].shape[1]).tolist()
+    word_ids = tokenizer(words, is_split_into_words=True, truncation=True, max_length=max_length).word_ids()
 
-    word_labels = []
-    used = set()
+    # first subtoken per word -> label
+    results = []
+    seen = set()
     for tok_i, w_i in enumerate(word_ids):
-        if w_i is None:
-            continue
-        if w_i in used:
+        if w_i is None or w_i in seen:
             continue
-        used.add(w_i)
-        word_labels.append((words[w_i], id2label[pred_ids[tok_i]]))
+        seen.add(w_i)
+        results.append((words[w_i], id2label[pred_ids[tok_i]]))
 
-    return word_labels
+    return results
 
 def find_label_vocab(vocabs):
     for i, v in enumerate(vocabs):
@@ -137,5 +135,6 @@ 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(sentence, model, label_vocab, device="cpu")
+pairs = predict_ner(sentence, model, tokenizer, id2label, device="cpu")
 print(pairs)