import gradio as gr
from transformers import GPTNeoModel, GPT2TokenizerFast

def generate_text(prompt, max_length, temperature, top_p, repetition_penalty):
    tokenizer = GPT2TokenizerFast.from_pretrained("JakeTurner616/Adonalsium-gpt-neo-1.3B")
    model = GPTNeoModel.from_pretrained("JakeTurner616/Adonalsium-gpt-neo-1.3B")
    
    # Check if tokenizer has a padding token, if not, add one
    if tokenizer.pad_token is None:
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
        model.resize_token_embeddings(len(tokenizer))  # Resize model embeddings to fit the new tokenizer

    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
    outputs = model.generate(
        input_ids=inputs["input_ids"],
        max_length=max_length,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        no_repeat_ngram_size=2
    )

    generated_texts = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
    return generated_texts

iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(lines=2, label="Input Prompt"),
        gr.Slider(minimum=10, maximum=300, step=10, value=100, label="Max Length"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.7, label="Temperature"),
        gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.9, label="Top P"),
        gr.Slider(minimum=1.0, maximum=2.0, step=0.1, value=1.1, label="Repetition Penalty")
    ],
    outputs="text",
)

iface.launch()