Update app.py

app.py

@@ -87,7 +87,6 @@ def tag_sentence(text):
     predictions = predict_ner_labels(model, tokenizer, text)
 
     # Obtenez les probabilités associées aux prédictions
-    # Vous devrez adapter cette partie en fonction de la sortie de votre modèle
     inputs = tokenizer(text, truncation=True, return_tensors="pt")
     outputs = model(**inputs)
     probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
@@ -95,9 +94,10 @@ def tag_sentence(text):
     # Calcul des probabilités que le tag prédit soit correct
     word_tags = []
     for i, tag in enumerate(predictions):
-        tag_id = id2tag.index(tag)
-        prob = np.round(probs[0, i, tag_id].item() * 100, 2)
-        word_tags.append((tokenizer.decode(inputs['input_ids'][0][i].item()), tag, prob))
+        tag_id = id2tag.get(tag, -1)  # Vérifiez si la clé existe, sinon utilisez -1 comme indice
+        if tag_id != -1:
+            prob = np.round(probs[0, i, tag_id].item() * 100, 2)
+            word_tags.append((tokenizer.decode(inputs['input_ids'][0][i].item()), tag, prob))
 
     # Créez un DataFrame avec les colonnes dans l'ordre spécifié
     df = pd.DataFrame(word_tags, columns=['word', 'tag', 'probability'])