Final

app.py

@@ -77,7 +77,7 @@ 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, tokenizer, id2label: dict, device="cpu", max_length=128):
+def predict_ner_nested(sentence, model, tokenizer, id2label, device="cpu", max_length=128):
     model.to(device)
     model.eval()
 
@@ -93,15 +93,37 @@ def predict_ner(sentence: str, model, tokenizer, id2label: dict, device="cpu", m
     enc = {k: v.to(device) for k, v in enc.items()}
 
     with torch.no_grad():
-        out = model(**enc)
-        logits = out.logits if hasattr(out, "logits") else out
-    
+        # 1) Get contextual token embeddings from the internal transformer
+        tf_out = model.transformer(
+            input_ids=enc["input_ids"],
+            attention_mask=enc.get("attention_mask", None)
+        )
+        x = tf_out.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)
+
+        # 3) Get logits
+        out = model(x, dummy_labels, segments_mask=None, get_sent_repr=False)
+
+        # 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
+
     pred_ids = logits.argmax(dim=-1)[0].tolist()
 
-    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()
+    # 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()
 
-    # first subtoken per word -> label
     results = []
     seen = set()
     for tok_i, w_i in enumerate(word_ids):
@@ -121,13 +143,16 @@ def find_label_vocab(vocabs):
     return None, None
 
 
-idx, label_vocab = find_label_vocab(label_vocab)
-print("label vocab index:", idx)
+label_vocab = label_vocab[0]  # the list loaded from pickle
 id2label = {i: s for i, s in enumerate(label_vocab.itos)}
 
+# idx, label_vocab = find_label_vocab(label_vocab)
+# print("label vocab index:", idx)
+# 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, tokenizer, id2label, device="cpu")
+pairs = predict_ner_nested(sentence, model, tokenizer, id2label, device="cpu")
 print(pairs)