Spaces:
Sleeping
Sleeping
File size: 4,763 Bytes
5e9edd3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
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]
|