import gradio as gr
import requests
import json
import os

API_URL = "https://api-inference.huggingface.co/models/openai-community/gpt2"
API_KEY = os.getenv('API_KEY')

if not API_KEY:
    raise ValueError("API_KEY environment variable not found. Please set it in your Secrets on Hugging Face.")

HEADERS = {"Authorization": f"Bearer {API_KEY}"}


def generate_text(prompt):
    payload = {
        "inputs": prompt,
        "parameters": {
            "max_length": 200,
            "do_sample": True
        }
    }
    response = requests.post(API_URL, headers=HEADERS, json=payload)
    
    if response.status_code == 200:
        return response.json()
    else:
        return f"Error: {response.status_code} - {response.text}"


description = "Enter your prompt and click Submit to see the model's response."
title = "Mistral Text Generation Interface"

api = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(label="Input Prompt", placeholder="Type something..."),
    outputs=gr.Textbox(label="Generated Text"),
    title=title,
    description=description
)


api.launch()