prelington commited on
Commit
a26baf9
·
verified ·
1 Parent(s): 468dc51

Create ProTalk_InstantChat.py

Browse files
Files changed (1) hide show
  1. ProTalk_InstantChat.py +31 -0
ProTalk_InstantChat.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoModelForCausalLM, AutoTokenizer
2
+ import torch
3
+
4
+ model_name = "microsoft/phi-2"
5
+ device = "cuda" if torch.cuda.is_available() else "cpu"
6
+
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16 if device=="cuda" else torch.float32).to(device)
9
+
10
+ system_prompt = "You are ProTalk, a professional AI assistant. Remember everything in this conversation. Be polite, witty, and professional."
11
+
12
+ chat_history = []
13
+
14
+ while True:
15
+ user_input = input("User: ")
16
+ if user_input.lower() == "exit":
17
+ break
18
+ chat_history.append(f"User: {user_input}")
19
+ prompt = system_prompt + "\n" + "\n".join(chat_history) + "\nProTalk:"
20
+ inputs = tokenizer(prompt, return_tensors="pt").to(device)
21
+ outputs = model.generate(
22
+ **inputs,
23
+ max_new_tokens=150,
24
+ do_sample=True,
25
+ temperature=0.7,
26
+ top_p=0.9,
27
+ repetition_penalty=1.2
28
+ )
29
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
30
+ print(f"ProTalk: {response}")
31
+ chat_history.append(f"ProTalk: {response}")