File size: 1,053 Bytes
468dc51 e9e5efd 468dc51 e9e5efd 468dc51 e9e5efd 468dc51 e9e5efd 468dc51 |
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 |
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_name = "./ProTalkModel.safetensors"
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("./")
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16 if device=="cuda" else torch.float32).to(device)
system_prompt = "You are ProTalk, a professional AI assistant. Remember everything in this conversation. Be polite, witty, and professional."
chat_history = []
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)
outputs = model.generate(**inputs, max_new_tokens=150, do_sample=True, temperature=0.7, top_p=0.9)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"ProTalk: {response}")
chat_history.append(f"ProTalk: {response}")
|