Spaces:
Sleeping
Sleeping
File size: 1,792 Bytes
8a0c27f b652e4e 8a0c27f b652e4e 19b2dc7 8a0c27f b652e4e 8a0c27f b652e4e 8a0c27f b652e4e 8a0c27f b652e4e 8a0c27f b652e4e 8a0c27f b652e4e 8a0c27f b652e4e 8a0c27f |
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 |
from typing import List, Sequence, Tuple
import faiss
import numpy as np
from core.vectorizer import Vectorizer
class PromptSearchEngine:
"""
The PromptSearchEngine is responsible for finding the most similar prompts to a given query
by leveraging vectorized representations of the prompts and a similarity search index.
"""
def __init__(self, prompts: Sequence[str]) -> None:
"""
Initialize the PromptSearchEngine with a list of prompts.
Args:
prompts (Sequence[str]): The sequence of raw corpus prompts to be indexed for similarity search.
"""
self.vectorizer = Vectorizer()
self.corpus_vectors = self.vectorizer.transform(prompts)
self.corpus = prompts
self.corpus_vectors = self.corpus_vectors / np.linalg.norm(
self.corpus_vectors, axis=1, keepdims=True
)
d = self.corpus_vectors.shape[1]
self.index = faiss.IndexFlatIP(d)
self.index.add(self.corpus_vectors.astype("float32"))
def most_similar(self, query: str, n: int = 5) -> List[Tuple[float, str]]:
"""
Find the most similar prompts to a given query.
Args:
query (str): The query prompt to search for similar prompts.
n (int, optional): The number of similar prompts to retrieve. Defaults to 5.
Returns:
List[Tuple[float, str]]: A list of tuples containing the similarity score and the corresponding prompt.
"""
query_vector = self.vectorizer.transform([query]).astype("float32")
query_vector = query_vector / np.linalg.norm(query_vector)
distances, indices = self.index.search(query_vector, n)
return [(distances[0][i], self.corpus[indices[0][i]]) for i in range(n)]
|