m

inference.py

@@ -68,7 +68,9 @@ def infer_araberta(text):
         predicted_ids = outputs['logits'][0].cpu().tolist()
 
     tokens = tokenizer.convert_ids_to_tokens(input_ids[0].cpu())
-    predicted_labels = [model.config.id2label.get(p, 'O') for p in predicted_ids]
+    # predicted_labels = [model.config.id2label.get(p, 'O') for p in predicted_ids]
+    predicted_labels = [model.id2label.get(p, 'O') for p in predicted_ids]
+
 
     clean_tokens = [t for t in tokens if t not in tokenizer.all_special_tokens]
     clean_labels = [l for t, l in zip(tokens, predicted_labels) if t not in tokenizer.all_special_tokens]