generate_submission.py

from datasets import load_dataset
import json
from extract_logic import run_groq_prompt

def answer_question(question):
    prompt = f"""
You are a general AI assistant. Answer the question below as simply and accurately as possible. No explanations, no units, no punctuation. Just the final answer.

Question: {question}
"""
    return run_groq_prompt(prompt).strip()

def generate_submission():
    print("🔁 Loading dataset...")
    dataset = load_dataset("gaia-benchmark/GAIA", "2023_level1", split="validation")
    subset = dataset.select(range(3))  # Keep it light for now

    results = []
    for example in subset:
        task_id = example["task_id"]
        question = example["Question"]
        print(f"📌 Answering task: {task_id}")
        try:
            answer = answer_question(question)
        except Exception as e:
            print(f"❌ Error: {e}")
            answer = "error"
        results.append({"task_id": task_id, "model_answer": answer})

    output_path = "/mnt/data/gaia_submission.jsonl"
    print(f"📁 Writing to {output_path}...")

    with open(output_path, "w", encoding="utf-8") as f:
        for r in results:
            f.write(json.dumps(r) + "\n")

    print("File exists?", os.path.exists("/mnt/data/gaia_submission.jsonl"))

    print(f"✅ Done! File at {output_path}")
    return output_path

if __name__ == "__main__":
    generate_submission()