Update app.py

app.py

@@ -54,14 +54,29 @@ def fill_masked(text: str, model_name: str, top_k: int):
     results = fill_mask(text, top_k=top_k)
 
     # Format results for display
-    formatted = []
-    for res in results:
-        formatted.append({
-            "sequence": res["sequence"],
-            "score": round(res["score"], 4),
-            "token": res["token_str"]
-        })
-    return formatted
+    res_dict = {}
+    if isinstance(results[0], list):
+        for mask_num, mas_results in enumerate(results):
+            formatted = []
+            for cand in mas_results:
+                formatted.append({
+                    "sequence": cand["sequence"],
+                    "score": round(cand["score"], 4),
+                    "token": cand["token_str"]
+                })
+            res_dict[mask_num] = formatted
+    elif
+        formatted = []
+            for cand in mas_results:
+                formatted.append({
+                    "sequence": cand["sequence"],
+                    "score": round(cand["score"], 4),
+                    "token": cand["token_str"]
+                })
+            res_dict[0] = formatted
+    else:
+        raise ValueError("The results need to be presented in a list or a list of lists")
+    return res_dict
 
 # Build the Gradio interface with a slider for top-k
 iface = gr.Interface(