import json

import ollama
from tqdm.notebook import tqdm

with open("../posts/cluster/cluster_reflect.ipynb") as file:
    notebook = json.load(file)

instruction = "Your job: \
- Improve grammar and language \
- fix errors \
- cut the clutter \
- but keep tone and voice \
- don't change markdown syntax, \
- rephrase for easier reading \
- limit the use of emoji \
- never cut jokes \
- only return the corrected text, no summaries of changes or list of improvements"

for cell in tqdm(notebook["cells"]):
    # print(cell["cell_type"])
    if cell["cell_type"] == "markdown":
        print("-" * 6)
        print(cell["source"])

        response = ollama.generate(
            model="phi",
            system=instruction,
            prompt=" ".join(cell["source"]),
            keep_alive="5m",
        )

        cell["source"].append("--- \n " + response["response"])

with open("../posts/cluster/cluster_reflect.ipynb", mode="w") as output_file:
    json.dump(notebook, output_file)