Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 from functools import partial
 import os
 
@@ -121,11 +122,11 @@ relik_out: RelikOutput = relik("Michael Jordan was one of the best players in th
 For more information, please refer to the [source code](https://github.com/SapienzaNLP/relik/).
 """
 
-wikipedia_retriever = GoldenRetriever("relik-ie/encoder-e5-base-v2-wikipedia")
-wikipedia_index = InMemoryDocumentIndex.from_pretrained("relik-ie/encoder-e5-base-v2-wikipedia-index", index_precision="bf16")
+wikipedia_retriever = GoldenRetriever("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 = GoldenRetriever("relik-ie/encoder-e5-small-v2-wikipedia-relations")
-wikidata_index = InMemoryDocumentIndex.from_pretrained("relik-ie/encoder-e5-small-v2-wikipedia-relations-index", index_precision="bf16")
+wikidata_retriever = GoldenRetriever("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")
 
 relik_available_models = [
     "relik-ie/relik-reader-small-cie-wikipedia",
@@ -137,6 +138,7 @@ relik_available_models = [
 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}},
@@ -158,12 +160,14 @@ relik_models = {
     "relik-ie/relik-relation-extraction-large-wikipedia": Relik.from_pretrained(
         "relik-ie/relik-relation-extraction-large-wikipedia",
         index=wikidata_index,
+        device="cuda",
         retriever=wikidata_retriever,
         device="cuda",
     ),
     "relik-ie/relik-entity-linking-large-robust": Relik.from_pretrained(
         "relik-ie/relik-entity-linking-large-robust",
         index=wikipedia_index,
+        device="cuda",
         retriever=wikipedia_retriever,
         reader_kwargs={"dataset_kwargs": {"use_nme": True}},
     ),
@@ -276,7 +280,7 @@ def generate_graph(spans, response, colors, dict_ents, bgcolor="#111827", font_c
     allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
     allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
 
-
+@spaces.GPU
 def text_analysis(Text, Model, Relation_Threshold, Window_Size, Window_Stride):
     global loaded_model
     if Model is None: