import os
import gradio as gr
from transformers import pipeline
from huggingface_hub import login

# Get the Hugging Face token from environment variables
HF_TOKEN = os.getenv('HF')

if not HF_TOKEN:
    raise ValueError("The HF environment variable is not set. Please set it to your Hugging Face token.")

# Authenticate with Hugging Face and save the token to the Git credentials helper
login(HF_TOKEN, add_to_git_credential=True)

# Create the pipeline for text generation using the specified model
pipe = pipeline("text-generation", model="distilbert/distilgpt2", token=HF_TOKEN)

def generate(text):
    try:
        # Generate the response using the pipeline
        responses = pipe(text, max_length=1024, num_return_sequences=1)
        response_text = responses[0]['generated_text']
        return response_text if response_text else "No valid response generated."
    
    except Exception as e:
        return str(e)

iface = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
    outputs="text",
    title="Chuunibyou Text Generator",
    description="Transform text into an elaborate and formal style with a nobleman tone.",
    live=False
)

def launch_custom_interface():
    iface.launch()
    with gr.TabbedInterface(fn=generate, inputs=gr.Textbox(lines=2, placeholder="Enter text here..."), outputs=gr.HTML(label="Output")) as ti:
        ti.add(custom_html)

if __name__ == "__main__":
    launch_custom_interface()