import os
import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
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)

# Load the model and tokenizer using the Hugging Face token
model_name = "google/gemma-7b-it"
tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name, token=HF_TOKEN)

def generate(text):
    try:
        # Tokenize the input text
        inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
        
        # Generate the response
        outputs = model.generate(
            inputs["input_ids"],
            max_length=1024,
            num_beams=5,
            early_stopping=True,
        )
        
        # Decode the output tokens to text
        response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        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()