Update app.py

app.py

@@ -79,14 +79,32 @@ def predict_ner_labels(model, tokenizer, sentence):
 
 id2tag = {0: 'O', 1: 'B-LOC', 2: 'B-PER', 3: 'I-PER', 4: 'B-ORG', 5: 'I-DATE', 6: 'B-DATE', 7: 'I-ORG', 8: 'I-LOC'}
 
+
 def tag_sentence(text):
     trainer, model, tokenizer = load_model()
+
+    # Utilisez votre modèle pour prédire les tags
     predictions = predict_ner_labels(model, tokenizer, text)
-    df = pd.DataFrame(predictions, columns=['tag'])
-    df['word'] = text.split()
-    df['probability'] = 100.0  # Vous pouvez ajuster cette valeur selon vos besoins
+
+    # 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)
+
+    # 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))
+
+    # Créez un DataFrame avec les colonnes dans l'ordre spécifié
+    df = pd.DataFrame(word_tags, columns=['word', 'tag', 'probability'])
+
     return df
 
+
 st.title("📘 Named Entity Recognition Wolof")
 
 with st.form(key='my_form'):