Final Updates

app.py

@@ -121,6 +121,19 @@ def predict_ner(sentence: str, model, id2label: dict, device="cpu"):
 
     return word_labels
 
+def find_label_vocab(vocabs):
+    for i, v in enumerate(vocabs):
+        if hasattr(v, "itos"):
+            itos = v.itos
+            if isinstance(itos, (list, tuple)) and any(x in itos for x in ["O", "B-PER", "I-PER"]):
+                return i, v
+    return None, None
+
+
+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)}