Spaces:
Sleeping
Sleeping
File size: 1,124 Bytes
a195930 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 | from datasets import load_dataset
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
print("Loading GAIA dataset...")
dataset = load_dataset(
"gaia-benchmark/GAIA",
"2023_level1",
split="validation"
)
print("Loading embedding model...")
embedder = SentenceTransformer(
"sentence-transformers/all-MiniLM-L6-v2"
)
questions = dataset["Question"]
embeddings = embedder.encode(
questions,
convert_to_numpy=True,
show_progress_bar=True
)
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)
print(f"Indexed {len(questions)} questions.")
def search_examples(query, k=3):
query_embedding = embedder.encode(
[query],
convert_to_numpy=True
)
distances, indices = index.search(
query_embedding,
k
)
examples = []
for idx in indices[0]:
row = dataset[int(idx)]
examples.append({
"question": row["Question"],
"answer": row.get("Final answer", ""),
"task_id": row["task_id"]
})
return examples |