update app

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import requests
 
 tokenizer = AutoTokenizer.from_pretrained("impresso-project/nel-hipe-multilingual")
 model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -9,21 +10,46 @@ model = AutoModelForSeq2SeqLM.from_pretrained(
 print("Model loaded successfully!")
 
 
+def get_wikipedia_title(qid, language="en"):
+    url = f"https://www.wikidata.org/w/api.php"
+    params = {
+        "action": "wbgetentities",
+        "format": "json",
+        "ids": qid,
+        "props": "sitelinks/urls",
+        "sitefilter": f"{language}wiki",
+    }
+
+    response = requests.get(url, params=params)
+    data = response.json()
+
+    try:
+        title = data["entities"][qid]["sitelinks"][f"{language}wiki"]["title"]
+        url = data["entities"][qid]["sitelinks"][f"{language}wiki"]["url"]
+        return title, url
+    except KeyError:
+        return "NIL", "None"
+
+
 def disambiguate_sentence(sentence):
-    results = []
     entities = []
-    for sentence in [sentence]:
-        outputs = model.generate(
-            **tokenizer([sentence], return_tensors="pt"),
-            num_beams=5,
-            num_return_sequences=5,
-            max_new_tokens=30,
-        )
-        decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-        results.append(decoded)
-        entities.append({"label": decoded[0]})
-
-        print(f"Decoded: {decoded}")
+    # Generate model outputs for the sentence
+    outputs = model.generate(
+        **tokenizer([sentence], return_tensors="pt"),
+        num_beams=5,
+        num_return_sequences=5,
+        max_new_tokens=30,
+    )
+    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    qid = decoded[0].split()[-1]  # Assuming QID is the last token in the output
+
+    # Get Wikipedia title and URL
+    title, url = get_wikipedia_title(qid)
+
+    entity_info = f"QID: {qid}, Title: {title}, URL: {url}"
+    entities.append(entity_info)
+
+    print(f"Entities: {entities}")
     return {"text": sentence, "entities": entities}
 
 
@@ -36,7 +62,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="Linked Entities")
+    output_entities = gr.Textbox(label="Linked Entities")
 
     # Interface definition
     interface = gr.Interface(