Create gaia_loader.py

gaia_loader.py

@@ -0,0 +1,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