Upload models

.gitattributes

@@ -34,3 +34,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index/documents.json filter=lfs diff=lfs merge=lfs -text
+models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index_filtered/documents.json filter=lfs diff=lfs merge=lfs -text
+frequency_blink.txt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,6 +1,6 @@
 ---
 title: Relik
-emoji: 📚
+emoji: 🤖
 colorFrom: red
 colorTo: yellow
 sdk: streamlit

app.py

@@ -181,7 +181,7 @@ def run_client():
 
     relik = Relik(
         question_encoder="/home/user/app/models/relik-retriever-small-aida-blink-pretrain-omniencoder/question_encoder",
-        document_index="/home/user/app/models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index",
+        document_index="/home/user/app/models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index_filtered",
         reader="/home/user/app/models/relik-reader-aida-deberta-small",
         top_k=100,
         window_size=32,

examples/explore_faiss.md

@@ -0,0 +1,8 @@
+# table to store results
+
+| Index          | nprobe | Recall | Time  |
+|----------------|--------|--------|-------|
+| Flat           |    1   | 98.7   | 38.64 |
+| IVFx,Flat      |    1   | 42.5   | 23.46 |
+| IVFx,Flat      |    14  | 88.5   | 133   |
+| IVFx_HNSW,Flat |    1   | 88.5   | 133   |

examples/explore_faiss.py

@@ -0,0 +1,163 @@
+import argparse
+import json
+import logging
+import os
+from pathlib import Path
+import time
+from typing import Union
+
+import torch
+import tqdm
+
+from relik.retriever import GoldenRetriever
+from relik.common.log import get_logger
+from relik.retriever.common.model_inputs import ModelInputs
+from relik.retriever.data.base.datasets import BaseDataset
+from relik.retriever.indexers.base import BaseDocumentIndex
+from relik.retriever.indexers.faiss import FaissDocumentIndex
+
+logger = get_logger(level=logging.INFO)
+
+
+def compute_retriever_stats(dataset) -> None:
+    correct, total = 0, 0
+    for sample in dataset:
+        window_candidates = sample["window_candidates"]
+        window_candidates = [c.replace("_", " ").lower() for c in window_candidates]
+
+        for ss, se, label in sample["window_labels"]:
+            if label == "--NME--":
+                continue
+            if label.replace("_", " ").lower() in window_candidates:
+                correct += 1
+            total += 1
+
+    recall = correct / total
+    print("Recall:", recall)
+
+
+@torch.no_grad()
+def add_candidates(
+    retriever_name_or_path: Union[str, os.PathLike],
+    document_index_name_or_path: Union[str, os.PathLike],
+    input_path: Union[str, os.PathLike],
+    batch_size: int = 128,
+    num_workers: int = 4,
+    index_type: str = "Flat",
+    nprobe: int = 1,
+    device: str = "cpu",
+    precision: str = "fp32",
+    topics: bool = False,
+):
+    document_index = BaseDocumentIndex.from_pretrained(
+        document_index_name_or_path,
+        # config_kwargs={
+        #     "_target_": "relik.retriever.indexers.faiss.FaissDocumentIndex",
+        #     "index_type": index_type,
+        #     "nprobe": nprobe,
+        # },
+        device=device,
+        precision=precision,
+    )
+
+    retriever = GoldenRetriever(
+        question_encoder=retriever_name_or_path,
+        document_index=document_index,
+        device=device,
+        precision=precision,
+        index_device=device,
+        index_precision=precision,
+    )
+    retriever.eval()
+
+    logger.info(f"Loading from {input_path}")
+    with open(input_path) as f:
+        samples = [json.loads(line) for line in f.readlines()]
+
+    topics = topics and "doc_topic" in samples[0]
+
+    # get tokenizer
+    tokenizer = retriever.question_tokenizer
+    collate_fn = lambda batch: ModelInputs(
+        tokenizer(
+            [b["text"] for b in batch],
+            text_pair=[b["doc_topic"] for b in batch] if topics else None,
+            padding=True,
+            return_tensors="pt",
+            truncation=True,
+        )
+    )
+    logger.info(f"Creating dataloader with batch size {batch_size}")
+    dataloader = torch.utils.data.DataLoader(
+        BaseDataset(name="passage", data=samples),
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=False,
+        collate_fn=collate_fn,
+    )
+
+    # we also dump the candidates to a file after a while
+    retrieved_accumulator = []
+    with torch.inference_mode():
+        start = time.time()
+        num_completed_docs = 0
+
+        for documents_batch in tqdm.tqdm(dataloader):
+            retrieve_kwargs = {
+                **documents_batch,
+                "k": 100,
+                "precision": precision,
+            }
+            batch_out = retriever.retrieve(**retrieve_kwargs)
+            retrieved_accumulator.extend(batch_out)
+
+        end = time.time()
+
+        output_data = []
+        # get the correct document from the original dataset
+        # the dataloader is not shuffled, so we can just count the number of
+        # documents we have seen so far
+        for sample, retrieved in zip(
+            samples[
+                num_completed_docs : num_completed_docs + len(retrieved_accumulator)
+            ],
+            retrieved_accumulator,
+        ):
+            candidate_titles = [c.label.split(" <def>", 1)[0] for c in retrieved]
+            sample["window_candidates"] = candidate_titles
+            sample["window_candidates_scores"] = [c.score for c in retrieved]
+            output_data.append(sample)
+
+        # for sample in output_data:
+        #     f_out.write(json.dumps(sample) + "\n")
+
+        num_completed_docs += len(retrieved_accumulator)
+        retrieved_accumulator = []
+
+    compute_retriever_stats(output_data)
+    print(f"Retrieval took {end - start:.2f} seconds")
+
+
+if __name__ == "__main__":
+    # arg_parser = argparse.ArgumentParser()
+    # arg_parser.add_argument("--retriever_name_or_path", type=str, required=True)
+    # arg_parser.add_argument("--document_index_name_or_path", type=str, required=True)
+    # arg_parser.add_argument("--input_path", type=str, required=True)
+    # arg_parser.add_argument("--output_path", type=str, required=True)
+    # arg_parser.add_argument("--batch_size", type=int, default=128)
+    # arg_parser.add_argument("--device", type=str, default="cuda")
+    # arg_parser.add_argument("--index_device", type=str, default="cpu")
+    # arg_parser.add_argument("--precision", type=str, default="fp32")
+
+    # add_candidates(**vars(arg_parser.parse_args()))
+    add_candidates(
+        "/root/relik-spaces/models/relik-retriever-small-aida-blink-pretrain-omniencoder/question_encoder",
+        "/root/relik-spaces/models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index_filtered",
+        "/root/relik-spaces/data/reader/aida/testa_windowed.jsonl",
+        # index_type="HNSW32",
+        # index_type="IVF1024,PQ8",
+        # nprobe=1,
+        topics=True,
+        device="cuda",
+    )

frequency_blink.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63bdea194b5c27d8c35547a205c42b4bc2e8933a47f179bc63256cf12a3bd448
+size 95579105

models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index_filtered/config.yaml

@@ -0,0 +1,8 @@
+_target_: relik.retriever.indexers.inmemory.InMemoryDocumentIndex
+documents:
+  _target_: relik.retriever.data.labels.Labels
+embeddings:
+  _target_: torch.Tensor
+name_or_dir: null
+device: cpu
+precision: null

models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index_filtered/documents.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:486ef055dcc484ddd9d445cfc2bac1e2a7c133d79492610de49b72630bd6ce8f
+size 719452975

models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index_filtered/embeddings.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee144610bf744e96091f4f295d350806173703d0960a964444a1c13b248a5c0d
+size 1537987243

relik/inference/annotator.py

@@ -4,6 +4,7 @@ from typing import Any, Callable, Dict, Optional, Union
 
 import hydra
 from omegaconf import OmegaConf
+from relik.retriever.indexers.faiss import FaissDocumentIndex
 from relik.retriever.pytorch_modules.hf import GoldenRetrieverModel
 from rich.pretty import pprint
 
@@ -395,10 +396,15 @@ class Relik:
 def main():
     from pprint import pprint
 
+    document_index = FaissDocumentIndex.from_pretrained(
+        "/root/relik-spaces/models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index",
+        config_kwargs={"_target_": "relik.retriever.indexers.faiss.FaissDocumentIndex", "index_type": "IVFx,Flat"},
+    )
+
     relik = Relik(
-        question_encoder="riccorl/relik-retriever-aida-blink-pretrain-omniencoder",
-        document_index="riccorl/index-relik-retriever-aida-blink-pretrain-omniencoder",
-        reader="riccorl/relik-reader-aida-deberta-small",
+        question_encoder="/root/relik-spaces/models/relik-retriever-small-aida-blink-pretrain-omniencoder/question_encoder",
+        document_index=document_index,
+        reader="/root/relik-spaces/models/relik-reader-aida-deberta-small",
         device="cuda",
         precision=16,
         top_k=100,

relik/retriever/__init__.py

@@ -0,0 +1 @@
+from relik.retriever.pytorch_modules.model import GoldenRetriever

relik/retriever/indexers/base.py

@@ -79,6 +79,17 @@ class BaseDocumentIndex:
         self.embeddings = embeddings
         self.name_or_dir = name_or_dir
 
+    def __iter__(self):
+        # make this class iterable
+        for i in range(len(self)):
+            yield self[i]
+
+    def __len__(self):
+        return self.documents.get_label_size()
+
+    def __getitem__(self, index):
+        return self.get_passage_from_index(index)
+
     @property
     def config(self) -> Dict[str, Any]:
         """
@@ -261,6 +272,7 @@ class BaseDocumentIndex:
         config_file_name: Optional[str] = None,
         document_file_name: Optional[str] = None,
         embedding_file_name: Optional[str] = None,
+        config_kwargs: Optional[Dict[str, Any]] = None,
         *args,
         **kwargs,
     ) -> "BaseDocumentIndex":
@@ -285,6 +297,9 @@ class BaseDocumentIndex:
             )
 
         config = OmegaConf.load(config_path)
+        # override the config with the kwargs
+        if config_kwargs is not None:
+            config = OmegaConf.merge(config, OmegaConf.create(config_kwargs))
         pprint(OmegaConf.to_container(config), console=console_logger, expand_all=True)
 
         # load the documents

relik/retriever/indexers/faiss.py

@@ -6,8 +6,9 @@ from dataclasses import dataclass
 from typing import Callable, List, Optional, Union
 
 import numpy
+import psutil
 import torch
-from pytorch_modules import RetrievedSample
+from relik.retriever.pytorch_modules import RetrievedSample
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
@@ -44,6 +45,7 @@ class FaissDocumentIndex(BaseDocumentIndex):
         embeddings: Optional[Union[torch.Tensor, numpy.ndarray]] = None,
         index=None,
         index_type: str = "Flat",
+        nprobe: int = 1,
         metric: int = faiss.METRIC_INNER_PRODUCT,
         normalize: bool = False,
         device: str = "cpu",
@@ -60,6 +62,8 @@ class FaissDocumentIndex(BaseDocumentIndex):
                     "The number of documents and embeddings must be the same."
                 )
 
+        faiss.omp_set_num_threads(psutil.cpu_count(logical=False))
+
         # device to store the embeddings
         self.device = device
 
@@ -83,6 +87,7 @@ class FaissDocumentIndex(BaseDocumentIndex):
                 self.embeddings = self._build_faiss_index(
                     embeddings=embeddings,
                     index_type=index_type,
+                    nprobe=nprobe,
                     normalize=normalize,
                     metric=metric,
                 )
@@ -91,6 +96,7 @@ class FaissDocumentIndex(BaseDocumentIndex):
         self,
         embeddings: Optional[Union[torch.Tensor, numpy.ndarray]],
         index_type: str,
+        nprobe: int,
         normalize: bool,
         metric: int,
     ):
@@ -103,11 +109,15 @@ class FaissDocumentIndex(BaseDocumentIndex):
         if self.normalize:
             index_type = f"L2norm,{index_type}"
         faiss_vector_size = embeddings.shape[1]
-        if self.device == "cpu":
-            index_type = index_type.replace("x,", "x_HNSW32,")
-        index_type = index_type.replace(
-            "x", str(math.ceil(math.sqrt(faiss_vector_size)) * 4)
-        )
+        # if self.device == "cpu":
+        #     index_type = index_type.replace("x,", "x_HNSW32,")
+        # nlist = math.ceil(math.sqrt(faiss_vector_size)) * 4
+        # # nlist = 8
+        # index_type = index_type.replace(
+        #     "x", str(nlist)
+        # )
+        # print("Current nlist:", nlist)
+        print("Current index:", index_type)
         self.embeddings = faiss.index_factory(faiss_vector_size, index_type, metric)
 
         # convert to GPU
@@ -121,12 +131,24 @@ class FaissDocumentIndex(BaseDocumentIndex):
                 embeddings.cpu() if isinstance(embeddings, torch.Tensor) else embeddings
             )
 
+        self.embeddings.hnsw.efConstruction = 20
         # convert to float32 if embeddings is a torch.Tensor and is float16
         if isinstance(embeddings, torch.Tensor) and embeddings.dtype == torch.float16:
             embeddings = embeddings.float()
 
+        logger.info("Training the index.")
+        self.embeddings.train(embeddings)
+        
+        logger.info("Adding the embeddings to the index.")
         self.embeddings.add(embeddings)
 
+        self.embeddings.nprobe = nprobe
+
+        # self.embeddings.hnsw.efSearch
+        self.embeddings.hnsw.efSearch = 256
+
+        # self.embeddings.k_factor = 10
+
         # save parameters for saving/loading
         self.index_type = index_type
         self.metric = metric
@@ -277,6 +299,7 @@ class FaissDocumentIndex(BaseDocumentIndex):
     @torch.no_grad()
     @torch.inference_mode()
     def search(self, query: torch.Tensor, k: int = 1) -> list[list[RetrievedSample]]:
+
         k = min(k, self.embeddings.ntotal)
 
         if self.normalize:
@@ -292,7 +315,7 @@ class FaissDocumentIndex(BaseDocumentIndex):
         batch_scores: List[List[float]] = retriever_out[0].detach().cpu().tolist()
         # Retrieve the passages corresponding to the indices
         batch_passages = [
-            [self.documents.get_label_from_index(i) for i in indices]
+            [self.documents.get_label_from_index(i) for i in indices if i != -1]
             for indices in batch_top_k
         ]
         # build the output object

relik/retriever/indexers/inmemory.py

@@ -67,6 +67,18 @@ class InMemoryDocumentIndex(BaseDocumentIndex):
                     f"Converting to {PRECISION_MAP[precision]}."
                 )
                 self.embeddings = self.embeddings.to(PRECISION_MAP[precision])
+        else:
+            if (
+                device == "cpu"
+                and self.embeddings is not None
+                and self.embeddings.dtype != torch.float32
+            ):
+                logger.info(
+                    "Index vectors are of type {}. Converting to float32.".format(
+                        self.embeddings.dtype
+                    )
+                )
+                self.embeddings = self.embeddings.to(PRECISION_MAP[32])
         # move the embeddings to the desired device
         if self.embeddings is not None and not self.embeddings.device == device:
             self.embeddings = self.embeddings.to(device)

requirements.txt

@@ -1,6 +1,6 @@
 #------- Core dependencies -------
 torch>=2.0
-transformers[sentencepiece]>=4.34,<4.35
+transformers[sentencepiece]>=4.33,<4.34
 rich>=13.0.0,<14.0.0
 scikit-learn
 overrides

scripts/blink_freq.py

@@ -0,0 +1,19 @@
+from collections import Counter
+import json
+
+from tqdm import tqdm
+
+if __name__ == "__main__":
+    counter = Counter()
+
+    with open("/media/data/EL/blink/train.alby-format.jsonl") as f_in:
+        for line in tqdm(f_in):
+            sample = json.loads(line)
+            for ss, se, label in sample["doc_annotations"]:
+                if label == "--NME--":
+                    continue
+                counter.update([label])
+
+    with open("frequency_blink.txt", "w") as f_out:
+        for k, v in counter.most_common():
+            f_out.write(f"{k}\t{v}\n")

scripts/filter_docs.py

@@ -0,0 +1,54 @@
+from collections import Counter
+import json
+import torch
+
+from tqdm import tqdm
+from relik.retriever.data.labels import Labels
+
+from relik.retriever.indexers.inmemory import InMemoryDocumentIndex
+
+if __name__ == "__main__":
+    with open("frequency_blink.txt") as f_in:
+        frequencies = [l.strip().split("\t")[0] for l in f_in.readlines()]
+
+    frequencies = set(frequencies[:1_000_000])
+
+    with open(
+        "/root/golden-retriever-v2/data/dpr-like/el/definitions_only_data.txt"
+    ) as f_in:
+        for line in f_in:
+            title = line.strip().split(" <def>")[0].strip()
+            frequencies.add(title)
+
+    document_index = InMemoryDocumentIndex.from_pretrained(
+        "/root/relik-spaces/models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index",
+    )
+
+    new_doc_index = {}
+    new_embeddings = []
+
+    for i in range(document_index.documents.get_label_size()):
+        doc = document_index.documents.get_label_from_index(i)
+        title = doc.split(" <def>")[0].strip()
+        if title in frequencies:
+            new_doc_index[doc] = len(new_doc_index)
+            new_embeddings.append(document_index.embeddings[i])
+
+    print(len(new_doc_index))
+    print(len(new_embeddings))
+
+    new_embeddings = torch.stack(new_embeddings, dim=0)
+    new_embeddings = new_embeddings.to(torch.float16)
+
+    print(new_embeddings.shape)
+
+    new_label_index = Labels()
+    new_label_index.add_labels(new_doc_index)
+    new_document_index = InMemoryDocumentIndex(
+        documents=new_label_index,
+        embeddings=new_embeddings,
+    )
+
+    new_document_index.save_pretrained(
+        "/root/relik-spaces/models/relik-retriever-small-aida-blink-pretrain-omniencoder/document_index_filtered"
+    )