import streamlit as st
import requests
import os

# Set up Hugging Face API Token and URL
API_TOKEN = os.getenv('API_TOKEN', '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')  # Use environment variable
API_URL = 'https://api-inference.huggingface.co/models/gpt2'  # GPT-2 model URL

headers = {
    "Authorization": f"Bearer {API_TOKEN}"
}

# Function to query the Hugging Face model
def query(prompt):
    payload = {"inputs": prompt}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

# Streamlit app interface
st.title("AI Text Generation with Hugging Face")

# Input prompt from user
user_input = st.text_area("Enter a prompt:", "Once upon a time")

# Button to trigger text generation
if st.button("Generate Text"):
    # Call the Hugging Face model with the user input
    response = query(user_input)
    generated_text = response[0]['generated_text']
    st.write(generated_text)