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