bug fixes in prediction function

app.py

@@ -4,12 +4,10 @@ import faiss
 
 from transformers import pipeline
 
-import requests
-
 
 sample_text = """Europejscy astronomowie odkryli planetę
 pozasłoneczną pochodzącą spoza naszej galaktyki, czyli 
-[START_ENT] Drogi Mlecznej [END_ENT]. Obserwacji dokonali
+[unused0] Drogi Mlecznej [unused1]. Obserwacji dokonali
 2,2-metrowym teleskopem MPG/ESO."""
 
 
@@ -34,18 +32,25 @@ def load_model(model_name: str = "clarin-knext/entity-linking-encoder"):
     return pipe
 
 
-def predict(model, index, query: str = sample_text, top_k: int=3):
+model = load_model()
+index = load_index()
+
+
+def predict(query: str = sample_text, top_k: int=3):
     index_data, faiss_index = index
-    query = model(query)
+    # takes only the [CLS] embedding (for now)
+    query = model(query, return_tensors = "pt")[0][0].numpy().reshape(1, -1)  
     
     scores, indices = faiss_index.search(query, top_k)
-    results = [index_data[idx] for row in indices for idx in row]
+    scores, indices = scores.tolist(), indices.tolist()
     
-    return "\n".join(str(results))
-
-
-model = load_model()
-index = load_index()
+    results = [
+        (index_data[result[0]], result[1])
+        for output in zip(indices, scores)
+        for result in zip(*output)
+    ]
+    
+    return str(results)
 
 
 demo = gr.Interface(fn=predict, inputs=textbox, outputs="text").launch()