import gradio as gr
from transformers import pipeline

generator = pipeline('text-generation', model="Pankaj8922/gpt2-fs", pad_token_id=50256)

def generate_text(prompt, max_length, temperature, top_k, top_p):
    result = generator(
        prompt, 
        max_length=max_length, 
        temperature=temperature, 
        top_k=top_k, 
        top_p=top_p, 
        do_sample=True, 
        truncation=True
    )
    return result[0]['generated_text']

with gr.Blocks() as demo:
    gr.Markdown("# GPT-2 Text From Scratch with Custom Settings")
    
    prompt = gr.Textbox(label="Enter your prompt here")
    
    max_length = gr.Slider(label="Max Length", minimum=10, maximum=300, value=50, step=1)
    temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.7, step=0.1)
    top_k = gr.Slider(label="Top-K Sampling", minimum=0, maximum=100, value=50, step=1)
    top_p = gr.Slider(label="Top-P (Nucleus Sampling)", minimum=0.0, maximum=1.0, value=0.9, step=0.1)
    
    output = gr.Textbox(label="Generated Text", interactive=False)
    
    generate_button = gr.Button("Generate")

    generate_button.click(generate_text, inputs=[prompt, max_length, temperature, top_k, top_p], outputs=output)

demo.launch()