habi01's picture
Update app.py
786d074 verified
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
def chat_response(message, history):
try:
# Build conversation history string
chat_history_ids = None
for human_msg, bot_msg in history:
# Encode user message
user_input_ids = tokenizer.encode(
human_msg + tokenizer.eos_token,
return_tensors='pt'
)
# Encode bot response
bot_output_ids = tokenizer.encode(
bot_msg + tokenizer.eos_token,
return_tensors='pt'
)
# Build full conversation
if chat_history_ids is None:
chat_history_ids = torch.cat([user_input_ids, bot_output_ids], dim=-1)
else:
chat_history_ids = torch.cat([chat_history_ids, user_input_ids, bot_output_ids], dim=-1)
# Add new user message
new_user_input_ids = tokenizer.encode(
message + tokenizer.eos_token,
return_tensors='pt'
)
# Generate response
chat_history_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if chat_history_ids is not None else new_user_input_ids
# Generate bot response
bot_output_ids = model.generate(
chat_history_ids,
max_length=1000,
pad_token_id=tokenizer.eos_token_id,
no_repeat_ngram_size=3,
do_sample=True,
top_k=100,
top_p=0.7,
temperature=0.8
)
# Extract only the bot's response (remove history)
response = tokenizer.decode(
bot_output_ids[:, chat_history_ids.shape[-1]:][0],
skip_special_tokens=True
)
return response
except Exception as e:
return f"Error: {str(e)}"
# Create chat interface
demo = gr.ChatInterface(
chat_response,
title="DialoGPT Chatbot",
examples=["Hello!", "What's AI?", "Tell me a joke"],
type="messages"
)
if __name__ == "__main__":
demo.launch()