#!/usr/bin/env python3
"""Interactive chat with Agent 1.2e. Run: python3 chat.py"""
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

MODEL = Path(__file__).parent / "model"
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16)
streamer = TextStreamer(tok, skip_prompt=True, skip_special_tokens=True)

history = []
print("Agent 1.2e — type 'exit' to quit, 'reset' to clear history.\n")
while True:
    try:
        user = input("you> ").strip()
    except (EOFError, KeyboardInterrupt):
        break
    if not user:
        continue
    if user == "exit":
        break
    if user == "reset":
        history = []
        print("(history cleared)")
        continue
    history.append({"role": "user", "content": user})
    enc = tok.apply_chat_template(history, add_generation_prompt=True,
                                  return_tensors="pt", return_dict=True)
    print("agent> ", end="", flush=True)
    # repetition_penalty matters here: merging with the base model weakens
    # the instruct model's stopping behavior, so it loops without it
    out = model.generate(**enc, max_new_tokens=512, do_sample=False,
                         repetition_penalty=1.15, streamer=streamer)
    reply = tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True)
    history.append({"role": "assistant", "content": reply})