Using huggingface's datasets library as key part of the pipeline about txtai HOT 4 CLOSED

class DenseInformationRetriever:

    _EMBEDDING_COL_NAME = "embeddings"

    def __init__(
        self,
        documents: nlp.Dataset,
        doc_vectorizer: Vectorizer,
        query_vectorizer: Vectorizer,
        batch_size: int = 512,
        cache_file: Optional[str] = None,
        string_factory: Optional[str] = "Flat",
        train_size: Optional[int] = None,
        min_train_pct: Optional[float] = 0.4,
    ):
        self.docs = documents
        self.batch_size = batch_size
        self.query_vectorizer = query_vectorizer
        self.docs = self.docs.map(
            lambda examples: {self._EMBEDDING_COL_NAME: doc_vectorizer(examples)},
            batched=True,
            batch_size=batch_size,
            cache_file_name=cache_file,
        )
        if train_size:
            train_size = min(train_size, int(min_train_pct * len(self.docs)))
        self.docs.add_faiss_index(
            column=self._EMBEDDING_COL_NAME,
            string_factory=string_factory,
            train_size=train_size,
        )

    def search(
        self, queries: nlp.Dataset, k: int = 10
    ) -> List[Tuple[List[float], Dict[str, List]]]:
        query_vectorizer = self.query_vectorizer
        embedding_col_name = self._EMBEDDING_COL_NAME

        queries = queries.map(
            lambda examples: {embedding_col_name: query_vectorizer(examples)},
            batched=True,
            batch_size=self.batch_size,
        )
        embeddings = queries[self._EMBEDDING_COL_NAME]
        results = []
        for i in tqdm.tqdm(range(0, len(embeddings), 32)):
            batch_scores, batch_retrieved = self.docs.get_nearest_examples_batch(
                self._EMBEDDING_COL_NAME,
                np.array(embeddings[i : i + 32], dtype=np.float32),
                k=k,
            )
            results.extend(list(zip(batch_scores, batch_retrieved)))
        return results

    def vector_search(
        self, vector: List[float], k: int = 10
    ) -> Tuple[List[float], Dict[str, List]]:
        scores, retrieved_examples = self.docs.get_nearest_examples(
            self._EMBEDDING_COL_NAME, np.array(vector, dtype=np.float32), k=k
        )
        return scores, retrieved_examples