# app.py (second app by claude)
import gradio as gr
import torch
from transformers import PreTrainedTokenizerFast
from .gpt_model_code import load_model_n_tokenizer, generate
# Load model and tokenizer once at startup
model, tokenizer = load_model_n_tokenizer()
model.eval()

def generate_text(prompt, max_new_tokens, top_k, top_p, temperature, repetition_penalty, penalize_len_below):
    device = next(model.parameters()).device
    
    # Convert top_k to None if using top_p
    if top_p > 0:
        top_k = None
    else:
        top_p = None
    
    with torch.no_grad():
        output_text = generate(  # function uses `with torch.no_grad()` internally already
            model=model,
            prompt=prompt,
            tokenizer=tokenizer,
            max_new_tokens=max_new_tokens,
            top_p=top_p,# top p sampling is prefered over top k if top_p != None
            top_k=top_k,
            temperature=0.7,
            repetition_penalty=repetition_penalty,  # New parameter: Repetition penalty factor
            penalize_len_below=penalize_len_below  # New parameter: Minimum content length for penalizing EOT token.
        )
        
        return output_text

# Create Gradio interface
with gr.Blocks(title="Nepali GPT-2 Text Generator") as interface:
    gr.Markdown("# Nepali GPT-2 Text Generator")
    gr.Markdown("Enter Nepali text to generate content using the custom GPT-2 model.")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="Enter Nepali text here...")
            max_tokens = gr.Slider(minimum=1, maximum=512, value=50, step=1, label="Max New Tokens")
            
            with gr.Row():
                with gr.Column():
                    top_k = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top K (set to 0 to use Top P)")
                    temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
                with gr.Column():
                    top_p = gr.Slider(minimum=0, maximum=1.0, value=0, step=0.05, label="Top P (set above 0 to use instead of Top K)")
                    repetition_penalty = gr.Slider(minimum=1.0, maximum=2.0, value=1.2, step=0.1, label="Repetition Penalty")
            
            min_length = gr.Slider(minimum=1, maximum=200, value=50, step=1, label="Minimum Length Penalty")
            generate_btn = gr.Button("Generate Text")
        
        with gr.Column():
            output = gr.Textbox(label="Generated Text", lines=10)
    
    # Add examples if you have any
    gr.Examples(
        examples=[
            ["रामले भात खायो", 50, 50, 0, 0.7, 1.2, 50],
            ["नेपाल एउटा", 100, 0, 0.9, 0.8, 1.2, 100],
        ],
        inputs=[prompt, max_tokens, top_k, top_p, temperature, repetition_penalty, min_length],
        outputs=output,
        fn=generate_text,
        cache_examples=True,
    )
    
    generate_btn.click(
        fn=generate_text,
        inputs=[prompt, max_tokens, top_p, top_k, temperature, repetition_penalty, min_length],
        outputs=output
    )


'''

'''

interface.launch()