harman / app.py
Ha4378422's picture
Create app.py
93f5dbe verified
raw
history blame contribute delete
970 Bytes
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load pre-trained model and tokenizer
model_name = "microsoft/DialoGPT-medium"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Function to generate a response
def generate_response(prompt):
inputs = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors="pt")
reply_ids = model.generate(inputs, max_length=1000, pad_token_id=tokenizer.eos_token_id)
reply = tokenizer.decode(reply_ids[:, inputs.shape[-1]:][0], skip_special_tokens=True)
return reply
# Simple chat loop
if __name__ == "__main__":
print("Chatbot: Hello! How can I help you today?")
while True:
user_input = input("You: ")
if user_input.lower() in ["exit", "quit", "bye"]:
print("Chatbot: Goodbye!")
break
response = generate_response(user_input)
print(f"Chatbot: {response}")