Update app.py

app.py

@@ -5,23 +5,15 @@ import spacy
 from spacy.tokens import Doc, Span
 from relik import Relik
 from relik.inference.data.objects import TaskType, RelikOutput
-from relik.retriever.indexers.inmemory import InMemoryDocumentIndex
 from pyvis.network import Network
 
 # RELIK Models Setup
-wikipedia_retriever = Relik.from_pretrained("relik-ie/encoder-e5-base-v2-wikipedia", device="cuda")
-wikipedia_index = InMemoryDocumentIndex.from_pretrained("relik-ie/encoder-e5-base-v2-wikipedia-index", index_precision="bf16", device="cuda")
-wikidata_retriever = Relik.from_pretrained("relik-ie/encoder-e5-small-v2-wikipedia-relations", device="cuda")
-wikidata_index = InMemoryDocumentIndex.from_pretrained("relik-ie/encoder-e5-small-v2-wikipedia-relations-index", index_precision="bf16", device="cuda")
+def setup_relik_model(model_name: str, device: str):
+    return Relik.from_pretrained(model_name, device=device)
 
 relik_models = {
-    "sapienzanlp/relik-entity-linking-large": Relik.from_pretrained(
-        "sapienzanlp/relik-entity-linking-large", device="cuda", index=wikipedia_index, retriever=wikipedia_retriever,
-        reader_kwargs={"dataset_kwargs": {"use_nme": True}}
-    ),
-    "relik-ie/relik-relation-extraction-small": Relik.from_pretrained(
-        "relik-ie/relik-relation-extraction-small", index=wikidata_index, device="cuda", retriever=wikidata_retriever
-    )
+    "sapienzanlp/relik-entity-linking-large": setup_relik_model("sapienzanlp/relik-entity-linking-large", "cuda"),
+    "relik-ie/relik-relation-extraction-small": setup_relik_model("relik-ie/relik-relation-extraction-small", "cuda")
 }
 
 def get_span_annotations(response, doc):
@@ -84,4 +76,4 @@ with gr.Blocks(fill_height=True, css=css, theme=theme) as demo:
         allow_flagging="never"
     )
     if __name__ == "__main__":
-        demo.launch()
+        demo.launch()