File size: 769 Bytes
ada3cf7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
import pandas as pd
def search_kg(query, index_path, dataset_path, top_k=5):
index = faiss.read_index(index_path)
df = pd.read_json(dataset_path, lines=True)
model = SentenceTransformer('all-MiniLM-L6-v2')
query_embedding = model.encode([query], convert_to_tensor=True).cpu().numpy()
distances, indices = index.search(query_embedding, top_k)
results = []
for i in range(top_k):
idx = indices[0][i]
if 0 <= idx < len(df):
headline = df.iloc[idx]['headline']
description = df.iloc[idx]['short_description']
results.append(f"{headline}. {description}")
return " ".join(results)
|