|
from transformers import BertTokenizer, BertModel |
|
import torch |
|
import pandas as pd |
|
from sklearn.metrics.pairwise import cosine_similarity |
|
import numpy as np |
|
|
|
|
|
df = pd.read_parquet('embeddings.parquet') |
|
df = df.head(5) |
|
|
|
|
|
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') |
|
model = BertModel.from_pretrained('bert-base-uncased') |
|
|
|
def search_embeddings(query): |
|
|
|
inputs = tokenizer(query_text, return_tensors="pt", padding=True, truncation=True, max_length=512) |
|
|
|
|
|
with torch.no_grad(): |
|
outputs = model(**inputs) |
|
query_vector = outputs.last_hidden_state.mean(dim=1).squeeze().numpy() |
|
|
|
|
|
|
|
|
|
embedding_matrix = np.stack(df['embedding'].values) |
|
|
|
similarities = cosine_similarity([query_vector], embedding_matrix) |
|
|
|
|
|
top_indices = np.argsort(similarities[0])[::-1][:5] |
|
top_scores = similarities[0][top_indices] |
|
|
|
results = "" |
|
|
|
|
|
for index, score in zip(top_indices, top_scores): |
|
|
|
data = df.iloc[index] |
|
results += (f"Question: {data['text']} Answer: {data['answer']} ") |
|
return results |
|
|
|
|
|
|
|
|
|
|
|
|
|
|