update app

app.py

@@ -11,6 +11,7 @@ print("Model loaded successfully!")
 
 def disambiguate_sentence(sentence):
     results = []
+    entities = []
     for sentence in [sentence]:
         outputs = model.generate(
             **tokenizer([sentence], return_tensors="pt"),
@@ -20,8 +21,10 @@ def disambiguate_sentence(sentence):
         )
         decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
         results.append(decoded)
+        entities.append({"label": decoded[0]})
+
         print(f"Decoded: {decoded}")
-    return results
+    return {"text": sentence, "entities": entities}
 
 
 def nel_app_interface():
@@ -33,7 +36,7 @@ def nel_app_interface():
         "entity should be surrounded by `[START]` and `[END]`. // "
         "!Only one entity per sentence is supported at the moment!",
     )
-    output_entities = gr.HighlightedText(label="Extracted Entities")
+    output_entities = gr.HighlightedText(label="Linked Entities")
 
     # Interface definition
     interface = gr.Interface(