import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("sergeantson/GPT2_Large_Law")
model = AutoModelForCausalLM.from_pretrained("sergeantson/GPT2_Large_Law")

def generate_text(input_text, max_length, num_return_sequences, temperature, top_k, top_p):
    inputs = tokenizer(input_text, return_tensors="pt")
    output = model.generate(
        **inputs,
        max_length=max_length,
        num_return_sequences=num_return_sequences,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        no_repeat_ngram_size=2  # Prevents repeating n-grams
    )
    generated_texts = [tokenizer.decode(output[i], skip_special_tokens=True) for i in range(num_return_sequences)]
    return "\n\n".join(generated_texts)

# Set up the Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter a prompt here...", label="Input Text"),
        gr.Slider(minimum=10, maximum=200, value=50, step=1, label="Max Length"),
        gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of Return Sequences"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top-p")
    ],
    outputs="text",
    title="Legal Text Generator",
    description="Enter a prompt to generate legal text based on the input."
)

# Launch the interface
iface.launch(share=False)