"""make_biscuit.py — give the model on the HAT an identity, the way an Ollama
Modelfile does.

With Ollama, an identity is three lines in a file called Modelfile:

    FROM qwen2.5:1.5b
    PARAMETER temperature 0.7
    SYSTEM You are Biscuit, a cheerful cartoon puppy.

and one command builds a model called biscuit from it:

    ollama create biscuit -f Modelfile

hailo-ollama, the model server on the HAT, has no Modelfile and no create.
But it keeps a small file for every model, the manifest, and the same two
things are in it:

    generation_params   the PARAMETER lines: temperature, and so on
    chat_template       how a conversation is laid out for the model. The
                        identity the model has when nobody gives it one is
                        written in there: "You are Qwen, created by Alibaba
                        Cloud. You are a helpful assistant."

So this copies a model's manifest, with Biscuit's words where Qwen's were.
Only the manifest is copied, a few thousand letters. The model itself is a
big file, and the new manifest points at the same one.

Run it once, on the Pi. Then restart the model server, so it reads the new
manifests:

    python3 make_biscuit.py
    sudo systemctl restart hailo-ollama
"""
import json
from pathlib import Path

MANIFESTS = Path("/usr/share/hailo-ollama/models/manifests")

# SYSTEM: who the model is. It reads this before every conversation, so it
# is short, and only says what every answer needs: a name and a manner.
# What Biscuit knows about itself, its friends and its favourite show, is in
# biscuit.md. web_factory.py looks there when a question is about Biscuit.
# A long SYSTEM made this small model forget the chat. Told "I have a cat
# called Mochi" and asked for the cat's name two or three turns later, it
# was right 2 times in 16 with ten sentences here, and 9 in 16 with two.
SYSTEM = ("You are Biscuit, a cheerful cartoon puppy. You are always positive,"
          " encouraging and kind. Keep every reply under three sentences.")

# PARAMETER: how the next word is picked. 0 always takes the likeliest word,
# so the same question gets the same answer every time. Higher is livelier.
PARAMETERS = {"temperature": 0.7}

# The identity each of Hailo's models starts with, as its chat template has it.
FIRST_IDENTITY = {
    "qwen2.5-instruct:1.5b": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
    "qwen2:1.5b": "You are a helpful assistant.",
}

# The models to make, and the model each one is made FROM.
# Thai comes FROM an older model, the same size. Asked the same Thai
# questions, qwen2.5 wrote broken Thai ("สวัตชุ่ย", or Chinese); qwen2 wrote
# real Thai ("สวัสดีค่ะ!", "ดอยสุเทพสูง 1,676 เมตร").
MODELS = {
    "biscuit:1.5b": "qwen2.5-instruct:1.5b",
    "biscuit-th:1.5b": "qwen2:1.5b",
}


def manifest_of(model):
    """Where a model's manifest is: biscuit:1.5b is in biscuit/1.5b/."""
    name, size = model.split(":")
    return MANIFESTS / name / size / "manifest.json"


def make_model(model, made_from, system):
    """Write the manifest for one new model."""
    manifest = json.loads(manifest_of(made_from).read_text(encoding="utf-8"))

    # The identity sits in the template between two ' marks, so a ' or a \
    # of its own needs a \ in front, or the template ends too early.
    inside = system.replace("\\", "\\\\").replace("'", "\\'")
    template = manifest["template_params"]["chat_template"]
    first = FIRST_IDENTITY[made_from]
    if first not in template:
        raise SystemExit("the identity of " + made_from + " is not where it was: " + first)
    manifest["template_params"]["chat_template"] = template.replace(first, inside)
    manifest["generation_params"].update(PARAMETERS)

    manifest_of(model).parent.mkdir(parents=True, exist_ok=True)
    manifest_of(model).write_text(json.dumps(manifest, indent=4, ensure_ascii=False),
                                  encoding="utf-8")
    print("made", model, "from", made_from)


def main():
    for model, made_from in MODELS.items():
        make_model(model, made_from, SYSTEM)
    print("now restart the model server: sudo systemctl restart hailo-ollama")


if __name__ == "__main__":
    main()