File size: 1,622 Bytes
ad6e824 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | # 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()
|