Spaces:
Sleeping
Sleeping
File size: 2,299 Bytes
01f5415 |
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 |
from typing import Sequence, List, Tuple
from models.vectorizer import Vectorizer
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
class PromptSearchEngine:
def __init__(self, model_name='bert-base-nli-mean-tokens'):
self.model = SentenceTransformer(model_name)
# Initialize FAISS index with right number of dimensions
self.embedding_dimension = self.model.get_sentence_embedding_dimension()
self.index = faiss.IndexFlatL2(self.embedding_dimension) # Euclidian distance index - brute force for small datasets
self.prompts_track = [] # To keep track of original prompts for returning results
def add_prompts_to_vector_database(self, prompts):
embeddings = self.model.encode(prompts)
self.index.add(np.array(embeddings).astype('float32'))
self.prompts_track.extend(prompts)
def most_similar(self, query, top_k=5):
# Encode the query
query_embedding = self.model.encode([query]).astype('float32')
# Optimizovana pretraga ali moramo promeniti vrstu indeksa
distances, indices = self.index.search(query_embedding, top_k)
# Retrieve the corresponding prompts for the found indices
similar_prompts = [self.prompts_track[idx] for idx in indices[0]]
return similar_prompts, distances[0] # Return both the similar prompts and their distances
def cosine_similarity(query_vector: np.ndarray, corpus_vectors: np.ndarray) -> np.ndarray:
"""Compute the cosine similarity between a query vector and a set of corpus vectors.
Args: query_vector: The query vector to compare against the corpus vectors. corpus_vectors: The set of corpus vectors to compare against the query vector.
Returns: The cosine similarity between the query vector and the corpus vectors.
"""
similarities = {}
for index, vector in enumerate(corpus_vectors):
if np.linalg.norm(vector) == 0:
raise ValueError("One of the corpus vectors has zero norm.")
cos_similarity = np.dot(vector, query_vector) / (np.linalg.norm(vector) * np.linalg.norm(query_vector))
similarities[index] = cos_similarity
return similarities
|