import gradio as gr
import torch
from model import model, tokenizer, translate_sentence, translate_beam_search, device

def translate(text, decoding_strategy, beam_size):
    if decoding_strategy == "Greedy":
        return translate_sentence(text, model, tokenizer, device)
    else:
        return translate_beam_search(text, model, tokenizer, device, pad_token_id=tokenizer.token_to_id('[PAD]'), beam_size=beam_size)

with gr.Blocks() as demo:
    gr.Markdown("Hindi-English Translation")
    
    with gr.Row():
        inp = gr.Textbox(label="Input Sentence", placeholder="Enter sentence to translate", lines=2)
    
    with gr.Row():
        decoding = gr.Radio(["Greedy", "Beam Search"], value="Greedy", label="Decoding Strategy")
        beam = gr.Slider(minimum=2, maximum=10, step=1, value=4, label="Beam Size (for Beam Search)")
    
    out = gr.Textbox(label="Translation", lines=2)

    def handle_translate(text, decoding_strategy, beam_size):
        return translate(text, decoding_strategy, beam_size)

    btn = gr.Button("Translate")
    btn.click(fn=handle_translate, inputs=[inp, decoding, beam], outputs=out)

demo.launch()