Duck-Duck_GO / app.py
Joekd608's picture
app.py
492375f unverified
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
# HF_TOKEN = secret.HF_TOKEN
# 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']}")