import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from huggingface_hub import HfApi

# Get the latest model from your space
api = HfApi()
space_name = "umut-bozdag/humanizer_model"  # Replace with your actual space name
model_files = api.list_repo_files(space_name)
model_file = next(file for file in model_files if file.endswith('.bin'))
model_revision = api.get_repo_info(space_name).sha

# Load the model and tokenizer from the space
tokenizer = AutoTokenizer.from_pretrained(space_name, revision=model_revision)
model = AutoModelForSeq2SeqLM.from_pretrained(space_name, revision=model_revision)

def generate_text(input_text):
    # Preprocess input text
    input_text = input_text.strip()
    
    # Prepare input for the model
    input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=256, truncation=True)
    
    # Generate text with parameters matching your training setup
    outputs = model.generate(
        input_ids,
        max_length=256,
        num_return_sequences=1,
        no_repeat_ngram_size=2,
        top_k=30,
        top_p=0.9,
        temperature=0.7,
        do_sample=True,
        early_stopping=True,
        num_beams=4
    )
    
    # Decode and clean up the generated text
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text.strip()

iface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs=gr.Textbox(label="Generated Text"),
    title="Text Humanizer",
    description="Enter text to generate a more human-like version."
)

iface.launch()