Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
# Load GPT-2 model and tokenizer | |
MODEL_NAME = "distilgpt2" | |
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME) | |
# Conversation history | |
conversation_history = "" | |
def chat(user_input): | |
global conversation_history | |
# Add user input to conversation history | |
conversation_history += f"User: {user_input}\n" | |
# Encode the conversation and generate a response | |
inputs = tokenizer(conversation_history, return_tensors="pt", truncation=True, max_length=1000) | |
outputs = model.generate(**inputs, max_length=200, num_return_sequences=1) | |
# Decode and add bot response to history | |
response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
conversation_history += f"Bot: {response}\n" | |
return response | |
# Gradio interface | |
interface = gr.Interface( | |
fn=chat, | |
inputs=gr.Textbox(placeholder="Type your message here..."), | |
outputs="text", | |
title="GPT-2 Chatbot", | |
description="A simple chatbot powered by GPT-2" | |
) | |
if __name__ == "__main__": | |
interface.launch() | |