"""Style suggestions for a few sentences in different languages. Run it from
the model folder:

    python example.py

Needs `mlx-lm` (Apple silicon). The prompt lives in `protocol.py` and must
not be changed: the model was trained on exactly that text.
"""
import os
import time

from mlx_lm import generate, load

import protocol as P

HERE = os.path.dirname(os.path.abspath(__file__))

TEXTS = [
    ("en", "Due to the fact that the office is closed, we will send the documents on Monday."),
    ("en", "The plumber comes on Tuesday between 9 and 12."),
    ("it", "Al fine di risparmiare, abbiamo deciso di spegnere la caldaia di notte."),
    ("de", "Zum jetzigen Zeitpunkt ist die Werkstatt geschlossen."),
    ("fr", "Nous avons procédé à la vérification des compteurs ce matin."),
    ("es", "El curso de cocina tuvo una duración de seis semanas."),
]


def main():
    model, tokenizer = load(HERE)
    for lang, text in TEXTS:
        tokens = tokenizer.apply_chat_template(P.messages(text, P.LANGUAGE_NAMES[lang]),
                                               tokenize=True, add_generation_prompt=True)
        t0 = time.time()
        answer = generate(model, tokenizer, prompt=tokens, max_tokens=256, verbose=False)
        found = P.style_suggestions(answer, text)
        print("%s  %s" % (lang, text))
        for o, r in found:
            print("    «%s» → «%s»" % (o, r))
        if not found:
            print("    (nothing to change)")
        print("    %.2f s" % (time.time() - t0))


if __name__ == "__main__":
    main()