import spaces
import gradio as gr
from transformers import pipeline, GPT2TokenizerFast

#model_id = "alakxender/dv-wiki-gpt2"
model_id = "alakxender/dv-articles-gpt2"
tokenizer = GPT2TokenizerFast.from_pretrained(model_id, model_max_length=128)

generator = pipeline("text-generation", model=model_id, tokenizer=tokenizer, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id)

@spaces.GPU
def generate_text(prompt, max_length, temperature):
    try:
        generated = generator(
            prompt,
            max_length=max_length,
            #num_beams=10,
            #no_repeat_ngram_size=2,
            temperature=temperature,
            do_sample=True,
            repetition_penalty=1.4
        )
        return generated[0]['generated_text']
    except Exception as e:
        return f"Something went wrong, try again. Error: {str(e)}"

styles = """
.thaana textarea {
    font-size: 18px !important;
    font-family: 'MV_Faseyha', 'Faruma', 'A_Faruma', 'Noto Sans Thaana', 'MV Boli';
    line-height: 1.8 !important;
}
"""

def create_interface():
    with gr.Blocks(css=styles) as demo:
        gr.Markdown("# Dhivehi Text Generator (GPT-2, Wiki)")
        gr.Markdown(
            "This is a GPT-2 model trained from Dhivehi text data from wikipedia\n"
            "Enter some text and generate a new text, adjust the parameters to generate text."
        )
        gr.Markdown("""
            **Parameters:**
            
            - **Temperature**: Controls the creativity of the output.
              - Lower values (0.2) = More focused and predictable text
              - Higher values (0.8) = More diverse and creative text
            
            - **Maximum Length**: Controls the length of generated text.
              - Higher values generate longer, more detailed results
              - Note: Longer texts take more time to generate
        """)
        
        with gr.Row():
            input_temperature = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.7,
                step=0.1,
                label="Temperature",
            )
            input_max_length = gr.Slider(
                minimum=10,
                maximum=128,
                value=60,
                step=1,
                label="Maximum Length",
            )
        
        with gr.Row():
            with gr.Column(scale=1):
                input_prompt = gr.Textbox(
                    label="Enter dhivehi text prompt",
                    placeholder="ދިވެހިން",
                    lines=5,
                    rtl=True,
                    elem_classes="thaana"
                )
            
            with gr.Column(scale=1):
                output_text = gr.Textbox(
                    label="Generated Text",
                    lines=5,
                    interactive=True,
                    rtl=True,
                    elem_classes="thaana"
                )
        
        with gr.Row():
            generate_btn = gr.Button("Generate", variant="primary")
            clear_btn = gr.ClearButton([input_prompt, output_text])
        
        generate_btn.click(
            fn=generate_text,
            inputs=[input_prompt, input_max_length, input_temperature],
            outputs=output_text
        )
        
        gr.Examples(
            examples=[
                ["ދިވެހިރާއްޖެ"],
                ["އެމެރިކާ އިންތިޚާބު"],
                ["ސަލާމް"],
                ["ދުނިޔޭގެ ސިއްޙަތު ޖަމްޢިއްޔާ"],
                ["ޤަދީމީ ސަގާފަތް"],
                ["ޑިމޮކްރަސީ"]
            ],
            inputs=input_prompt
        )
    
    return demo            

if __name__ == "__main__":
    demo = create_interface()
    demo.queue().launch(show_api=False)