import pandas as pd
from transformers import pipeline
from prompts import prompts_masked

fill = pipeline("fill-mask", model="bert-base-uncased")

results = []
for prompt in prompts_masked:
    out = fill(prompt, top_k=10)   # top predictions
    results.append({
        "prompt": prompt,
        "top1": out[0]["token_str"].strip(),
        "top1_score": out[0]["score"],
        "top5": [x["token_str"].strip() for x in out[:5]]
    })

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