prelington/CodexTrouter
1
1!pip install -q transformers torch accelerate safetensors2 3from transformers import AutoModelForCausalLM, AutoTokenizer4import torch5 6model_name = "microsoft/phi-2"7 8device = "cuda" if torch.cuda.is_available() else "cpu"9 10tokenizer = AutoTokenizer.from_pretrained(model_name)11model = AutoModelForCausalLM.from_pretrained(12 model_name,13 torch_dtype=torch.float16 if device=="cuda" else torch.float32,14 device_map="auto" if device=="cuda" else None15)16 17system_prompt = "You are ProTalk, a professional AI assistant. Remember everything in this conversation. Be polite, witty, and professional."18chat_history = []19 20while True:21 user_input = input("User: ")22 if user_input.lower() == "exit":23 break24 chat_history.append(f"User: {user_input}")25 prompt = system_prompt + "\n" + "\n".join(chat_history) + "\nProTalk:"26 inputs = tokenizer(prompt, return_tensors="pt").to(device)27 outputs = model.generate(28 **inputs,29 max_new_tokens=150,30 do_sample=True,31 temperature=0.7,32 top_p=0.9,33 repetition_penalty=1.234 )35 response = tokenizer.decode(outputs[0], skip_special_tokens=True)36 print(f"ProTalk: {response}")37 chat_history.append(f"ProTalk: {response}")38 