Update app.py

app.py

@@ -68,12 +68,36 @@ input_text = st.text_area(
 tokenizer_roberta = AutoTokenizer.from_pretrained("Cicciokr/Roberta-Base-Latin-Uncased")
 model_roberta = AutoModelForMaskedLM.from_pretrained("Cicciokr/Roberta-Base-Latin-Uncased")
 fill_mask_roberta = pipeline("fill-mask", model=model_roberta, tokenizer=tokenizer_roberta)
+punctuation_marks = {".", ",", ";", ":", "!", "?"}
+
+def get_valid_predictions(sentence, max_attempts=3, top_k=5):
+    attempt = 0
+    filtered_predictions = []
+
+    while attempt < max_attempts:
+        predictions = fill_mask_roberta(sentence, top_k=top_k)
+        
+        # Filtra le predizioni rimuovendo la punteggiatura
+        filtered_predictions = [
+            pred for pred in predictions if pred["token_str"] not in punctuation_marks
+        ]
+        
+        # Se troviamo almeno una parola valida, interrompiamo il ciclo
+        if filtered_predictions:
+            break
+        
+        attempt += 1
+
+    return filtered_predictions
+
 
 # Se l'utente ha inserito (o selezionato) un testo
 if input_text:
     # Sostituiamo [MASK] con <mask> (lo tokenizer Roberta se lo aspetta così)
     input_text_roberta = input_text.replace("[MASK]", "<mask>")
-    predictions_roberta = fill_mask_roberta(input_text_roberta)
+
+    
+    predictions_roberta = get_valid_predictions(input_text_roberta)
     
     st.subheader("Risultati delle previsioni:")
     for pred in predictions_roberta: