AI / app.py
Liam2343's picture
Update app.py
0fd311c verified
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
# Load the pre-trained model
model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
chat_history_ids = None
# Function to handle the chat logic
def chat(user_input, history=[]):
global chat_history_ids
new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
if chat_history_ids is not None:
input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
else:
input_ids = new_input_ids
chat_history_ids = model.generate(
input_ids,
max_length=1000,
pad_token_id=tokenizer.eos_token_id,
temperature=0.7,
top_p=0.9,
no_repeat_ngram_size=3,
top_k=50,
do_sample=True,
)
response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
history.append((user_input, response))
return history, history
# Custom HTML header
custom_html = """
<div style="text-align:center; padding: 20px; background-color: #1e1e2f; color: white; border-radius: 12px;">
<h1>🤖 Smart AI Assistant</h1>
<p>Talk to DialoGPT and experience AI conversation in real time.</p>
</div>
"""
# Gradio interface setup
def launch_gradio_interface():
with gr.Interface(fn=chat,
inputs=[gr.Textbox(placeholder="Say something...", elem_id="user-input"), gr.State([])],
outputs=[gr.Chatbot(), gr.State()],
title="Smart AI Assistant",
live=True) as demo:
# Adding custom HTML
demo.add_component(gr.HTML(custom_html), row=0, col=0)
demo.launch()
launch_gradio_interface()