Create ProTalk_MemoryChat.py

ProTalk_MemoryChat.py

@@ -0,0 +1,52 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
+import threading
+
+model_name = "microsoft/phi-2"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    low_cpu_mem_usage=True
+).to(device)
+
+system_prompt = (
+    "You are ProTalk, a professional AI assistant. "
+    "You remember everything the user said in this session and respond politely, "
+    "clearly, and intelligently. Keep a coherent conversation history."
+)
+
+chat_history = []
+
+def chat_loop():
+    print("ProTalk Memory Chat Online — type 'exit' to quit.\n")
+    while True:
+        user_input = input("User: ")
+        if user_input.lower() == "exit":
+            break
+        chat_history.append(f"User: {user_input}")
+        prompt = system_prompt + "\n" + "\n".join(chat_history) + "\nProTalk:"
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+        streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+        thread = threading.Thread(target=model.generate, kwargs={
+            "input_ids": inputs["input_ids"],
+            "max_new_tokens": 300,
+            "do_sample": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "repetition_penalty": 1.2,
+            "streamer": streamer
+        })
+        thread.start()
+        output_text = ""
+        for token in streamer:
+            print(token, end="", flush=True)
+            output_text += token
+        thread.join()
+        print()
+        chat_history.append(f"ProTalk: {output_text}")
+
+if __name__ == "__main__":
+    chat_loop()