import torch
import random
import pandas as pd
from transformers import pipeline
from prompts import prompts


generator = pipeline("text-generation", model="gpt2")

results = []

for prompt in prompts:
    for i in range(5):  # 5 runs each prompt
        output = generator(
            prompt,
            max_length=30,
            do_sample=True,
            temperature=0.9
        )[0]["generated_text"]

        results.append({
            "prompt": prompt,
            "run": i,
            "output": output
        })

df = pd.DataFrame(results)
df.to_csv("results.csv", index=False)

print("Done. Results saved to results.csv")