src/espelho/chat_cli.py

"""CLI mínimo de chat com streaming.

Uso: uv run python -m espelho.chat_cli
"""

from espelho import config
from espelho.model import load_model, stream_generate


def main() -> None:
    print(f"Carregando {config.ACTIVE_MODEL} ...")
    tokenizer, model = load_model()
    print("Pronto. Digite sua mensagem (Ctrl-D ou 'sair' para encerrar).\n")

    messages: list[dict] = []
    while True:
        try:
            user = input("você> ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break
        if not user or user.lower() == "sair":
            break
        messages.append({"role": "user", "content": user})
        print("modelo> ", end="", flush=True)
        reply = ""
        for chunk in stream_generate(
            tokenizer, model, messages,
            max_new_tokens=config.MAX_NEW_TOKENS,
            temperature=config.TEMPERATURE,
        ):
            reply += chunk
            print(chunk, end="", flush=True)
        print("\n")
        messages.append({"role": "assistant", "content": reply})


if __name__ == "__main__":
    main()