Upload scripts/prompt_engine.py with huggingface_hub

scripts/prompt_engine.py

@@ -0,0 +1,104 @@
+import os
+import sys
+sys.path.append(sys.path[0].replace('scripts', ''))
+import pandas as pd
+import numpy as np
+
+from config.data_paths import VECTORDB_PATH
+
+from typing import Sequence, List, Tuple
+import faiss
+from sentence_transformers import SentenceTransformer
+
+
+class Vectorizer:
+    def __init__(self, model_name: str) -> None:
+        """
+        Initialize the vectorizer with a pre-trained embedding model.
+        Args:
+            model_name: The name of the pre-trained embedding model (compatible with sentence-transformers).
+        """
+        self.model = SentenceTransformer(model_name)
+
+    def transform(self, prompts: Sequence[str], build_index=False) -> np.ndarray:
+        """
+        Transform texts into numerical vectors using the specified model.
+        Args:
+            prompts: The sequence of raw corpus prompts.
+        Returns:
+            Vectorized prompts as a numpy array.
+        """
+        embeddings = self.model.encode(prompts, show_progress_bar=True)
+        embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True) # normalize embeddings
+        if build_index:
+            # self.embeddings=embeddings
+            if os.path.isfile(os.path.join(VECTORDB_PATH, 'prompts_index.faiss')):
+                print('Embeddings already stored in vector db')
+            else:
+                index = self._build_index(embeddings)
+                faiss.write_index(index, os.path.join(VECTORDB_PATH, 'prompts_index.faiss'))
+        else:
+            return embeddings
+
+    def _build_index(self, embeddings: np.ndarray) -> faiss.IndexFlatIP:
+        """
+        Build and return a FAISS index for the given embeddings.
+        Args:
+            embeddings: A numpy array of prompt embeddings.
+        Returns:
+            FAISS index for efficient similarity search.
+        """
+        index = faiss.IndexFlatIP(embeddings.shape[1])  # Cosine similarity (IP on normalized vectors)
+        index.add(embeddings)
+        return index
+
+def cosine_similarity(query_vector: np.ndarray, corpus_vectors: np.ndarray) -> np.ndarray:
+    """
+    Calculate cosine similarity between prompt vectors.
+    Args:
+        query_vector: Vectorized prompt query of shape (1, D).
+        corpus_vectors: Vectorized prompt corpus of shape (N, D).
+    Returns:
+        A vector of shape (N,) with values in range [-1, 1] where 1 is maximum similarity.
+    """
+    return np.dot(corpus_vectors, query_vector.T).flatten()
+
+class PromptSearchEngine:
+    def __init__(self, corpus: str, model_name: str = 'all-MiniLM-L6-v2', use_index=False) -> None:
+        """
+        Initialize search engine by vectorizing prompt corpus.
+        Vectorized prompt corpus should be used to find the top n most similar prompts.
+        Args:
+            corpus: Path to the parquet dataset with raw prompts.
+            model_name: The name of the pre-trained embedding model.
+        """
+        self.use_index=use_index
+        self.prompts=pd.read_parquet(corpus)['prompt'].to_list()
+        self.prompts=self.prompts# if use_index else np.random.choice(self.prompts, 1000, replace=False)
+        self.vectorizer = Vectorizer(model_name)
+        self.embeddings = self.vectorizer.transform(self.prompts, 
+                                  build_index=use_index) # build index initially for faster retrieval
+        if use_index:
+            self.index = faiss.read_index(os.path.join(VECTORDB_PATH, 'prompts_index.faiss'))
+
+    def most_similar(self, query: str, n: int = 5) -> List[Tuple[float, str]]:
+        """
+        Return top n most similar prompts from the corpus.
+        Input query prompt is vectorized using the Vectorizer. After that, use the cosine_similarity
+        function to get the top n most similar prompts from the corpus.
+        Args:
+            query: The raw query prompt input from the user.
+            n: The number of similar prompts to return from the corpus.
+        Returns:
+            The list of top n most similar prompts from the corpus along with similarity scores.
+            Note that returned prompts are verbatim.
+        """
+        query_vector = self.vectorizer.transform([query])
+        if self.use_index:
+            distances, indices = self.index.search(query_vector, n)
+            results = [{'prompt': self.prompts[idx], 'score': distances[0][i]} for i, idx in enumerate(indices[0])]
+            return results
+        else:
+            similarities = cosine_similarity(query_vector, self.embeddings)
+            top_indices = np.argsort(-similarities)[:n]  # Sort in descending order
+            return [{'prompt': self.prompts[i], 'score': similarities[i]} for i in top_indices]