Create chatbot_local.py

chatbot_local.py

@@ -0,0 +1,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:]