prelington/CodexTrouter

1from transformers import AutoModelForCausalLM, AutoTokenizer2import torch3 4model_name = "microsoft/phi-2"5device = "cuda" if torch.cuda.is_available() else "cpu"6 7tokenizer = AutoTokenizer.from_pretrained(model_name)8model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16 if device=="cuda" else torch.float32).to(device)9 10system_prompt = "You are ProTalk, a professional AI assistant. Remember everything in this conversation. Be polite, witty, and professional."11 12chat_history = []13 14while True:15    user_input = input("User: ")16    if user_input.lower() == "exit":17        break18    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.228    )29    response = tokenizer.decode(outputs[0], skip_special_tokens=True)30    print(f"ProTalk: {response}")31    chat_history.append(f"ProTalk: {response}")32