#!/usr/bin/env python
"""
Runnable example for abhiram3040/simplewords-dictation-cleanup-v2.

    pip install mlx-lm huggingface_hub
    python example.py                 # run the built-in demo cases
    python example.py "your um raw text"

Two things are load-bearing and must not be changed:
  1. SYSTEM is frozen. It must match training BYTE-FOR-BYTE. It ships in the repo
     as system_v2.txt and is read from there rather than retyped.
  2. Decoding is GREEDY (temperature 0) and thinking is DISABLED. Sampling or a
     re-enabled <think> block will degrade or corrupt the output.
"""
import sys
from pathlib import Path

from huggingface_hub import snapshot_download
from mlx_lm import load, generate
from mlx_lm.sample_utils import make_sampler

REPO = "abhiram3040/simplewords-dictation-cleanup-v2"

DEMO = [
    "let's meet on tuesday wait no friday at noon",
    "um so like can you uh send the report to the team by tomorrow",
    "red one no the blue one actually the green one",
    "send it tuesday i mean before noon",
    "make a list of milk, eggs, bread and coffee",
    "tell sarah the macbook shipped",
]


def main() -> None:
    path = snapshot_download(REPO)
    # The frozen prompt ships with the weights -- read it, never retype it.
    system = (Path(path) / "system_v2.txt").read_text().strip()
    model, tok = load(path)
    sampler = make_sampler(temp=0.0)                      # GREEDY

    def clean(raw: str) -> str:
        prompt = tok.apply_chat_template(
            [{"role": "user", "content": f"{system}\n\n{raw}"}],
            add_generation_prompt=True,
            enable_thinking=False,                        # no reasoning trace
            tokenize=False,
        )
        return generate(model, tok, prompt=prompt,
                        max_tokens=512, sampler=sampler).strip()

    for raw in (sys.argv[1:] or DEMO):
        print(f"raw   : {raw}")
        print(f"clean : {clean(raw)}\n")


if __name__ == "__main__":
    main()