# example_chat.py — minimal streaming chat with OpenModel-Labs/CSM-1-Preview
# pip install -U torch transformers accelerate flash-linear-attention huggingface_hub
import sys, threading, torch
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, AutoModelForImageTextToText, TextIteratorStreamer

REPO = "OpenModel-Labs/CSM-1-Preview"
system = open(hf_hub_download(REPO, "system_prompt.txt"), encoding="utf-8").read()

tok = AutoTokenizer.from_pretrained(REPO)
model = AutoModelForImageTextToText.from_pretrained(REPO, dtype=torch.bfloat16, device_map="auto")
model.eval()

messages = [{"role": "system", "content": system}]
print("Cyber Safety Model - type a question, Ctrl+C to quit.")
while True:
    try:
        q = input("\nyou > ").strip()
    except (EOFError, KeyboardInterrupt):
        break
    if not q:
        continue
    messages.append({"role": "user", "content": q})
    prompt = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False, enable_thinking=False)
    inputs = tok(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
    streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
    kw = dict(**inputs, streamer=streamer, max_new_tokens=1024, do_sample=True,
              temperature=0.7, top_p=0.9, top_k=20)
    threading.Thread(target=lambda: model.generate(**kw), daemon=True).start()
    reply = ""
    for chunk in streamer:
        reply += chunk
        sys.stdout.write(chunk)
        sys.stdout.flush()
    messages.append({"role": "assistant", "content": reply.strip()})
    print()