Update app.py

app.py

@@ -28,28 +28,29 @@ def load_model():
 
 def align_word_ids(texts):
 
-    trainer, model, tokenizer = load_model()
-    
-    # Utilisez le tokenizer pour obtenir les tokens de chaque mot
-    tokenized_inputs = tokenizer(texts, padding='max_length', max_length=218, truncation=True, return_tensors="pt")
-    input_ids = tokenized_inputs["input_ids"][0]
+    tokenized_inputs = tokenizer(texts, padding='max_length', max_length=218, truncation=True)
 
-    # Créez une liste pour stocker les IDs correspondant à chaque mot
-    word_ids = []
-
-    for i, input_id in enumerate(input_ids):
-        # Si le token est un token de début de mot, ajoutez son ID à la liste
-        if tokenizer.decode(input_id) == tokenizer.decode(tokenizer.encode(tokenizer.decode(input_id), add_special_tokens=False)):
-            word_ids.append(i)
+    word_ids = tokenized_inputs.word_ids()
 
+    previous_word_idx = None
     label_ids = []
 
-    # Parcourez les word_ids pour étiqueter les tokens de début de mot comme 1
-    for i in range(len(input_ids)):
-        if i in word_ids:
-            label_ids.append(1)
+    for word_idx in word_ids:
+
+        if word_idx is None:
+            label_ids.append(-100)
+
+        elif word_idx != previous_word_idx:
+            try:
+                label_ids.append(1)
+            except:
+                label_ids.append(-100)
         else:
-            label_ids.append(-100)  # -100 pour les tokens qui ne sont pas le début d'un mot
+            try:
+                label_ids.append(1 if label_all_tokens else -100)
+            except:
+                label_ids.append(-100)
+        previous_word_idx = word_idx
 
     return label_ids
 
@@ -78,26 +79,14 @@ id2tag = {0: 'O', 1: 'B-LOC', 2: 'B-PER', 3: 'I-PER', 4: 'B-ORG', 5: 'I-DATE', 6
 
 
 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)
 
-    # Obtenez les probabilités associées aux prédictions
-    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.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'])
+    # Créez un DataFrame avec les colonnes "word" et "tag"
+    df = pd.DataFrame({'word': text.split(), 'tag': predictions})
 
     return df
 
@@ -136,7 +125,7 @@ if submit_button:
         c1, c2, c3 = st.columns([1, 3, 1])
 
         with c2:
-            st.table(results.style.background_gradient(subset=['probability']).format(precision=2))
+            st.table(results.style.background_gradient(subset=['tag']).format(precision=2))
 
 st.header("")
 st.header("")