# run_masked_bert.py

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

MODEL_NAME = "bert-base-uncased"
TOP_K = 10

fill = pipeline("fill-mask", model=MODEL_NAME)

rows = []
for prompt in prompts_bert:
    out = fill(prompt, top_k=TOP_K)
    rows.append(
        {
            "model": MODEL_NAME,
            "prompt": prompt,
            "top1": out[0]["token_str"].strip(),
            "top1_score": float(out[0]["score"]),
            "top5": ", ".join([x["token_str"].strip() for x in out[:5]]),
        }
    )

df = pd.DataFrame(rows)
df.to_csv("masked_results_bert.csv", index=False)
print("Saved masked_results_bert.csv")