File size: 3,214 Bytes
37abe8d | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import torch
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
# ==================================================
# CONFIG
# ==================================================
# Either the HF hub id (auto-downloads + caches on first run):
MODEL_PATH = "Qwen/Qwen2.5-0.5B-Instruct"
# ...or a local folder if you already downloaded it yourself
# (e.g. via snapshot_download or git lfs clone), e.g.:
# MODEL_PATH = "./Qwen2.5-0.5B-Instruct"
SYSTEM_PROMPT = "You are a helpful assistant."
MAX_NEW_TOKENS = 512
HISTORY_TURNS = 6 # number of past user/assistant exchanges kept as context
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
# ==================================================
# LOAD MODEL
# ==================================================
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, torch_dtype=dtype)
model.to(device)
model.eval()
print("Device:", device)
# ==================================================
# PROMPT BUILDING
# ==================================================
def build_prompt(history, user_input):
messages = [{"role": "system", "content": SYSTEM_PROMPT}]
for role, text in history:
messages.append({"role": role, "content": text})
messages.append({"role": "user", "content": user_input})
return tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
# ==================================================
# GENERATION (streamed)
# ==================================================
def generate_reply(history, user_input):
prompt = build_prompt(history, user_input)
inputs = tokenizer(prompt, return_tensors="pt").to(device)
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt=True,
skip_special_tokens=True,
)
generation_kwargs = dict(
**inputs,
streamer=streamer,
max_new_tokens=MAX_NEW_TOKENS,
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.1,
pad_token_id=tokenizer.eos_token_id,
)
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
full_text = ""
for token in streamer:
print(token, end="", flush=True)
full_text += token
thread.join()
print()
return full_text
# ==================================================
# CHAT
# ==================================================
print("\n" + "=" * 60)
print("Qwen2.5-0.5B-Instruct Chatbot")
print("Type exit to quit")
print("=" * 60)
history = [] # list of (role, text) tuples, most recent last
while True:
user_input = input("\nUser: ")
if user_input.lower() == "exit":
break
if not user_input.strip():
continue
print("\nAssistant:")
reply = generate_reply(history, user_input)
history.append(("user", user_input))
history.append(("assistant", reply))
history = history[-HISTORY_TURNS * 2:] |