import json
from transformers import pipeline

# Load a basic model (replace with your preferred one)
agent = pipeline("text-generation", model="gpt2", max_new_tokens=100)

# Example GAIA task subset (mocked)
subset_tasks = [
    {"task_id": "gaia_0001", "input": "What is the capital of France?"},
    {"task_id": "gaia_0002", "input": "List 3 crops commonly grown in Saskatchewan."},
    {"task_id": "gaia_0003", "input": "Explain why climate change affects agriculture."}
]

# Create predictions
submission = []

for task in subset_tasks:
    prompt = task["input"]
    response = agent(prompt)[0]["generated_text"]

    submission.append({
        "task_id": task["task_id"],
        "model_answer": response.strip(),
        "reasoning_trace": f"Used GPT-2 to generate answer to: '{prompt}'"
    })

# Save to JSON Lines format
with open("submission.jsonl", "w", encoding="utf-8") as f:
    for entry in submission:
        f.write(json.dumps(entry) + "\n")

print("Submission file saved as submission.jsonl")