File size: 2,042 Bytes
799adc9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
# Load the model and tokenizer
model_name = "distilgpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Function to generate response and maintain chat history
def chat_function(user_input, history):
if history is None:
history = []
# Create prompt with history
prompt = "\n".join([f"User: {h[0]}\nAI: {h[1]}" for h in history] + [f"User: {user_input}"])
# Tokenize input
inputs = tokenizer(prompt, return_tensors="pt", padding=True)
# Generate response
outputs = model.generate(
inputs["input_ids"],
max_length=100,
num_return_sequences=1,
temperature=0.7,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
# Decode response
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
response = response[len(prompt):].strip() or "Hmm, I'm not sure what to say!"
# Update history
history.append((user_input, response))
return history, history
# Create Gradio interface
with gr.Blocks(title="Simple Chat App") as demo:
gr.Markdown("# Simple AI Chat App")
gr.Markdown("Chat with an AI powered by DistilGPT-2!")
# Chatbot component for displaying conversation
chatbot = gr.Chatbot(label="Conversation")
# Input box
user_input = gr.Textbox(label="Your message", placeholder="Type here...")
# Hidden state to maintain chat history
history = gr.State(value=[])
# Submit button
submit_btn = gr.Button("Send")
# Clear button
clear_btn = gr.Button("Clear Chat")
# Connect components
submit_btn.click(
fn=chat_function,
inputs=[user_input, history],
outputs=[chatbot, history]
)
clear_btn.click(
fn=lambda: ([], []),
inputs=None,
outputs=[chatbot, history]
)
# Launch the app
demo.launch() |