| 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) | |