GPT2 / app.py
Waseem7711's picture
Update app.py
144cc69 verified
# Import necessary libraries
import streamlit as st
from transformers import pipeline, set_seed
# Set title for the Streamlit app
st.title("GPT-2 Chatbot")
# Load the GPT-2 model using Hugging Face pipeline
@st.cache_resource
def load_gpt2_pipeline():
"""Load the GPT-2 text generation pipeline."""
generator = pipeline('text-generation', model='gpt2')
set_seed(42) # Set seed for reproducibility
return generator
# Function to generate text from user input
def generate_text(prompt, generator, max_length=50, num_return_sequences=1):
"""Generate text from a given prompt."""
return generator(prompt, max_length=max_length, num_return_sequences=num_return_sequences)
# Load the GPT-2 pipeline
generator = load_gpt2_pipeline()
# Create an input text box for user prompt
user_input = st.text_input("Enter your prompt:", "Hello, how are you?")
# Settings in sidebar for generation
st.sidebar.header("Settings")
max_length = st.sidebar.slider("Max Length", min_value=10, max_value=100, value=50)
num_return_sequences = st.sidebar.slider("Number of Responses", min_value=1, max_value=5, value=1)
# Generate response when the button is pressed
if st.button("Generate Response"):
with st.spinner("Generating response..."):
responses = generate_text(user_input, generator, max_length, num_return_sequences)
for i, response in enumerate(responses):
st.write(f"**Response {i+1}:** {response['generated_text']}")