| | from abc import ABC, abstractmethod |
| |
|
| | import pandas as pd |
| | import torch |
| | from datasets import load_from_disk |
| | from sentence_transformers import SentenceTransformer |
| |
|
| | |
| |
|
| |
|
| | class TextEmbedder(ABC): |
| | def __init__(self, model_name, paragraphs_path, device, load_existing_index=False): |
| | self.dataset = load_from_disk(paragraphs_path) |
| | self.model = self._load_model(model_name, device) |
| |
|
| | assert len(self.dataset) > 0, "The loaded dataset is empty !!" |
| |
|
| | if load_existing_index == True: |
| | self.dataset.load_faiss_index( |
| | "embeddings", f"{paragraphs_path}/index.faiss" |
| | ) |
| |
|
| | def generate_paragraphs_embedding(self): |
| | self.dataset = self.dataset.map( |
| | lambda x: {"embeddings": self._generate_embeddings(x["content"])} |
| | ) |
| |
|
| | def save_embeddings(self, output_path): |
| | self.dataset.add_faiss_index(column="embeddings") |
| | self.dataset.save_faiss_index("embeddings", f"{output_path}/index.faiss") |
| |
|
| | def retrieve_faiss(self, query: str, k_total: int, threshold: int): |
| | question_embedding = self._generate_embeddings(query) |
| | scores, samples = self.dataset.get_nearest_examples( |
| | "embeddings", question_embedding, k=k_total |
| | ) |
| | passages_df = pd.DataFrame(samples) |
| | passages_df["scores"] = scores / 100 |
| | passages_df = passages_df[passages_df["scores"] > threshold] |
| | passages_df = passages_df.sort_values(by=["scores"], ascending=False) |
| |
|
| | if len(passages_df) == 0: |
| | return [], [] |
| |
|
| | contents = passages_df["content"].tolist() |
| | meta = passages_df.drop(columns=["content"]).to_dict(orient="records") |
| | passages = [] |
| | for i in range(len(contents)): |
| | passages.append({"content": contents[i], "meta": meta[i]}) |
| | return passages, passages_df["scores"].values |
| |
|
| | def retrieve_elastic(self, query: str, k_total: int, threshold: int): |
| | raise NotImplementedError |
| |
|
| | @abstractmethod |
| | def _load_model(self, model_name: str, device: str): |
| | pass |
| |
|
| | @abstractmethod |
| | def _generate_embeddings(self, text: str): |
| | pass |
| |
|
| |
|
| | class SentenceTransformersTextEmbedder(TextEmbedder): |
| | def _load_model(self, model_name: str, device: str): |
| | model = SentenceTransformer(model_name) |
| | torch_device = torch.device(device) |
| | model.to(torch_device) |
| | return model |
| |
|
| | def _generate_embeddings(self, text: str): |
| | return self.model.encode(text) |
| |
|
| |
|
| | |
| | |
| | |
| | |
| |
|
| | |
| | |
| | |
| |
|
