Update

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
+from transformers import AutoTokenizer, AutoModel
 app = FastAPI()
 print("Version 2...")
 
@@ -13,8 +13,11 @@ print("Version 2...")
     # filename="tag_vocab.pkl"
 # )
 
-pretrained_path = "aubmindlab/bert-base-arabertv2"  # change if different in your training
+
+pretrained_path = "aubmindlab/bert-base-arabertv2"  # set to your training backbone
 tokenizer = AutoTokenizer.from_pretrained(pretrained_path)
+encoder = AutoModel.from_pretrained(pretrained_path)
+encoder.eval()
 
 checkpoint_path = hf_hub_download(
     repo_id="SinaLab/Nested",
@@ -77,9 +80,9 @@ ckpt = torch.load(checkpoint_path, map_location="cpu")
 model = load_model_from_checkpoint(model, ckpt, strict=False)
 # model.eval()
 
-def predict_ner_nested(sentence, model, tokenizer, id2label, device="cpu", max_length=128):
-    model.to(device)
-    model.eval()
+def predict_ner_with_external_encoder(sentence, tagger, encoder, tokenizer, id2label, device="cpu", max_length=128):
+    tagger.to(device).eval()
+    encoder.to(device).eval()
 
     words = sentence.split()
 
@@ -93,36 +96,26 @@ def predict_ner_nested(sentence, model, tokenizer, id2label, device="cpu", max_l
     enc = {k: v.to(device) for k, v in enc.items()}
 
     with torch.no_grad():
-        # 1) Get contextual token embeddings from the internal transformer
-        tf_out = model.transformer(
+        x = encoder(
             input_ids=enc["input_ids"],
             attention_mask=enc.get("attention_mask", None)
-        )
-        x = tf_out.last_hidden_state  # [1, seq_len, hidden]
+        ).last_hidden_state  # [1, seq_len, hidden]
 
-        # 2) Dummy labels (because forward requires labels)
-        seq_len = x.size(1)
-        ignore_idx = getattr(model, "label_ignore_idx", 0)
-        dummy_labels = torch.full((1, seq_len), ignore_idx, dtype=torch.long, device=device)
+        ignore_idx = getattr(tagger, "label_ignore_idx", 0)
+        dummy_labels = torch.full((1, x.size(1)), ignore_idx, dtype=torch.long, device=device)
 
-        # 3) Get logits
-        out = model(x, dummy_labels, segments_mask=None, get_sent_repr=False)
+        out = tagger(x, dummy_labels)
 
-        # Your forward may return logits or (loss, logits) or dict-like
         if isinstance(out, (tuple, list)):
             logits = out[-1]
         elif hasattr(out, "logits"):
             logits = out.logits
         else:
-            logits = out  # assume tensor
+            logits = out
 
     pred_ids = logits.argmax(dim=-1)[0].tolist()
 
-    # Map tokens back to words (first subtoken per word)
-    if hasattr(enc, "word_ids"):
-        word_ids = enc.word_ids(batch_index=0)
-    else:
-        word_ids = tokenizer(words, is_split_into_words=True, truncation=True, max_length=max_length).word_ids()
+    word_ids = tokenizer(words, is_split_into_words=True, truncation=True, max_length=max_length).word_ids()
 
     results = []
     seen = set()
@@ -134,6 +127,7 @@ def predict_ner_nested(sentence, model, tokenizer, id2label, device="cpu", max_l
 
     return results
 
+
 def find_label_vocab(vocabs):
     for i, v in enumerate(vocabs):
         if hasattr(v, "itos"):
@@ -143,6 +137,8 @@ def find_label_vocab(vocabs):
     return None, None
 
 
+
+
 label_vocab = label_vocab[0]  # the list loaded from pickle
 id2label = {i: s for i, s in enumerate(label_vocab.itos)}
 
@@ -154,5 +150,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_nested(sentence, model, tokenizer, id2label, device="cpu")
+pairs = predict_ner_with_external_encoder(sentence, model, encoder, tokenizer, id2label, device="cpu")
 print(pairs)