|
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) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
