Nexialog/CSRD_GPT

1from abc import ABC, abstractmethod2 3import pandas as pd4import torch5from datasets import load_from_disk6from sentence_transformers import SentenceTransformer7 8# from finbert_embedding.embedding import FinbertEmbedding9 10 11class TextEmbedder(ABC):12    def __init__(self, model_name, paragraphs_path, device, load_existing_index=False):13        """Initialize an instance of the TextEmbedder class.14        Args:15            model_name (str): The name of the SentenceTransformer model to be used for embeddings.16            paragraphs_path (str): The path to the dataset of paragraphs to be embedded.17            device (str): The target device to run the model ('cpu' or 'cuda').18            load_existing_index (bool): If True, load an existing Faiss index, if available.19        Returns:20            None21        """22        self.dataset = load_from_disk(paragraphs_path)23        self.model = self._load_model(model_name, device)24 25        assert len(self.dataset) > 0, "The loaded dataset is empty !!"26 27        if load_existing_index == True:28            self.dataset.load_faiss_index(29                "embeddings", f"{paragraphs_path}/index.faiss"30            )31 32    # Generate embeddings for each paragraph 33    def generate_paragraphs_embedding(self):34        """Generate embeddings for paragraphs in the dataset.35            This function computes embeddings for each paragraph's content in the dataset and adds36            the embeddings as a new column named "embeddings" to the dataset.37        Args:38            None39        Returns:40            None41        """42        self.dataset = self.dataset.map(43            lambda x: {"embeddings": self._generate_embeddings(x["content"])}44        )45 46    # Save embeddings47    def save_embeddings(self, output_path):48        """Save Faiss embeddings index to a specified output path.49        Args:50            output_path (str): The path to save the Faiss embeddings index.51        Returns:52            None53        """54        self.dataset.add_faiss_index(column="embeddings")55        self.dataset.save_faiss_index("embeddings", f"{output_path}/index.faiss")56 57    # Allows the search 58    def retrieve_faiss(self, query: str, k_total: int, threshold: int):59        """Retrieve passages using Faiss similarity search.60        Args:61            query (str): The query for which similar passages are to be retrieved.62            k_total (int): The total number of passages to retrieve.63            threshold (int): The minimum similarity score threshold for passages to be considered.64        Returns:65            Tuple[List[Dict[str, Union[str, Dict[str, Any]]], np.ndarray]]:66            A tuple containing:67            - List of dictionaries, each representing a passage with 'content' (str) and 'meta' (dict) fields.68            - Numpy array of similarity scores for the retrieved passages.69        """70        question_embedding = self._generate_embeddings(query)71        scores, samples = self.dataset.get_nearest_examples(72            "embeddings", question_embedding, k=k_total73        )74        passages_df = pd.DataFrame(samples)75        passages_df["scores"] = scores / 10076        passages_df = passages_df[passages_df["scores"] > threshold]77        passages_df = passages_df.sort_values(by=["scores"], ascending=False)78 79        if len(passages_df) == 0:80            return [], []81 82        contents = passages_df["content"].tolist()83        meta = passages_df.drop(columns=["content"]).to_dict(orient="records")84        passages = []85        for i in range(len(contents)):86            passages.append({"content": contents[i], "meta": meta[i]})87        return passages, passages_df["scores"].values88 89    def retrieve_elastic(self, query: str, k_total: int, threshold: int):90        raise NotImplementedError91 92    @abstractmethod93    def _load_model(self, model_name: str, device: str):94        pass95 96    @abstractmethod97    def _generate_embeddings(self, text: str):98        pass99 100 101class SentenceTransformersTextEmbedder(TextEmbedder):102    def _load_model(self, model_name: str, device: str):103        """Load a SentenceTransformer model onto the specified device.104        Args:105            model_name (str): The name of the SentenceTransformer model to be loaded.106            device (str): The target device to move the model to ('cpu' or 'cuda').107        Returns:108            SentenceTransformer: The loaded SentenceTransformer model placed on the specified device.109        """110        model = SentenceTransformer(model_name)111        torch_device = torch.device(device)112        model.to(torch_device)113        return model114 115    def _generate_embeddings(self, text: str):116        """Generate embeddings for a given text using the loaded model.117        Args:118            text (str): The input text for which embeddings are to be generated.119        Returns:120            np.ndarray: An array representing the embeddings of the input text.121        """122        return self.model.encode(text)123 124 125# class FinBertTextEmbedder(TextEmbedder):126#     def _load_model(self, model_name: str, device: str):127#         model = FinbertEmbedding(device=device)128#         return model129 130#     def _generate_embeddings(self, text: str):131#         output = self.model.sentence_vector(text)132#         return output.cpu().numpy()133