| |
| |
| 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() |
|
|
