Update app.py

app.py

@@ -14,6 +14,9 @@ from relik.inference.serve.frontend.utils import get_random_color
 from relik.retriever.pytorch_modules.model import GoldenRetriever
 from relik.retriever.indexers.inmemory import InMemoryDocumentIndex
 from relik.inference.data.objects import TaskType
+from relik.retriever.pytorch_modules import RetrievedSample
+from relik.retriever.indexers.document import Document, DocumentStore
+from relik.retriever.indexers.base import BaseDocumentIndex
 
 LOGO = """
 <div style="text-align: center; display: flex; flex-direction: column; align-items: center;">
@@ -122,17 +125,49 @@ 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/).
 """
 
+class GoldenSillyRetriever(GoldenRetriever):
+    def __init__(self, documents: List[str], *args, **kwargs):
+        self.documents = DocumentStore([Document(doc) for doc in documents])
+        self.document_index = BaseDocumentIndex(self.documents)
+    def retrieve(self,
+                text: Optional[Union[str, List[str]]] = None,
+                k: int = 100,
+                *args,
+                **kwargs,
+    ) -> List[List[RetrievedSample]]:
+        if isinstance(text, str):
+            text = [text]
+        elif text is None:
+            text = []
+        return [
+            [RetrievedSample(score=1.0, document=doc) for doc in self.documents[:k]]
+            for _ in text
+        ]
+    def index(self):
+        pass
+    def eval(self):
+        pass
+    def save_pretrained(self):
+        pass
+    def to(self, device):
+        pass
+
 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", device="cuda")
 wikidata_index = InMemoryDocumentIndex.from_pretrained("relik-ie/encoder-e5-small-v2-wikipedia-relations-index", index_precision="bf16", device="cuda")
 
+ner_type_retriever = GoldenSillyRetriever(
+                documents=['media', 'disease', 'miscellaneous', 'event', 'person', 'location', 'time', 'celestial', 'organization', 'concept']
+            )
+
 relik_available_models = [
     "relik-ie/relik-reader-small-cie-wikipedia",
     "sapienzanlp/relik-entity-linking-large",
     "relik-ie/relik-entity-linking-large-robust",
-    "relik-ie/relik-relation-extraction-large-wikipedia",
+    "relik-ie/relik-relation-extraction-small-wikipedia-ner",
+    "relik-ie/relik-relation-extraction-small-wikipedia",
 ]
 
 relik_models = {
@@ -157,8 +192,22 @@ relik_models = {
         },
         reader_kwargs={"dataset_kwargs": {"use_nme": True}},
     ),
-    "relik-ie/relik-relation-extraction-large-wikipedia": Relik.from_pretrained(
-        "relik-ie/relik-relation-extraction-large-wikipedia",
+    "relik-ie/relik-relation-extraction-small-wikipedia-ner": Relik.from_pretrained(
+        "relik-ie/relik-relation-extraction-small-wikipedia-ner",
+        index=wikidata_index,
+        device="cuda",
+        use_nme=True,
+        retriever={
+            TaskType.SPAN: ner_type_retriever,
+            TaskType.TRIPLET: wikidata_retriever,
+        },
+        index={
+            TaskType.SPAN: ner_type_retriever.document_index,
+            TaskType.TRIPLET: wikidata_index,
+        }
+    ),
+    "relik-ie/relik-relation-extraction-small-wikipedia": Relik.from_pretrained(
+        "relik-ie/relik-relation-extraction-small-wikipedia",
         index=wikidata_index,
         device="cuda",
         retriever=wikidata_retriever,