Spaces:
Running
Running
import streamlit as st | |
from transformers import pipeline | |
from huggingface_hub import login | |
import torch | |
import os | |
#### | |
# Set page configuration | |
st.set_page_config(page_title="Text GenAI Model", page_icon="🤖") | |
st.title("Text GenAI Model") | |
st.subheader("Answer Random Questions Using Hugging Face Models") | |
# Fetch Hugging Face token from Streamlit Secrets | |
access_token_read = st.secrets[HF_TOKEN] # Ensure this is set in your Streamlit Cloud Secrets | |
# Free up GPU memory (if using GPU) | |
torch.cuda.empty_cache() | |
# Set environment variable to avoid fragmentation | |
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" | |
# Login to Hugging Face Hub using the access token | |
login(token=access_token_read) | |
# Initialize the text generation pipeline with GPT-2 model | |
pipe = pipeline("text-generation", model="distilbert/distilgpt2") # Using CPU | |
# Input from the user | |
text = st.text_input("Ask a Random Question") | |
if text: | |
# Generate text based on the random question | |
response = pipe(f"Answer the question: {text}", max_length=150, num_return_sequences=1) | |
# Display the generated response | |
st.text(f"Answer: {response[0]['generated_text']}") | |