Team Ai
Modelpublic

prelington/CodexTrouter

sourceHugging Faceapache-2.0updated 1y agoView on Hugging Face
1likes
ProTalk_Colab.py38 linesDownload Raw Back to root
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