import gradio as gr
from transformers import pipeline

# Load the model as a translation pipeline
translator = pipeline("translation_en_to_he", model="guymorlan/LostInTranslation")
description_text="<div style='direction: rtl;'>הכניסו שם של סרט באנגלית כדי לשבש אותו ללא היכ.. כלומר לתרגם אותו לעברית 😈</div>"

explanation = """
<div dir='rtl'>
<p>נוצר על ידי <a href="mailto:guy.mor@mail.huji.ac.il">גיא מור-לן</a>.</p>

<h3>שיטות פיענוח (Decoding Methods)</h3>
<p>שיטת הפיענוח קובעת כיצד מחוללים טקסט באמצעות מודל גנרטיבי. יש הרבה שיטות, אבל כאן בחרתי בשתי שיטות נפוצות והגדרתי אותן באופן שיוביל לתרגומים מגוונים יחסית.</p>

<h4>חיפוש אלומה (Beam Search)</h4>
<p>Beam search is performed with 20 beams, 5 beam groups, and a diversity penalty of 10.0 to induce more diverse translations.</p>

<h4>דגימה (Sampling)</h4>
<p>Nucleus sampling is performed with a top-p of 0.5 and a temperature of 1.8. This method generates wilder translations, but also more nonsense.</p>
</div>
"""

def translate(title, method):
    if "אלומה" in method:
        # Beam search
        translations = translator(title, num_return_sequences=10, num_beams=20, num_beam_groups=5, diversity_penalty=10.0)
    else:
        # Sampling
        translations = translator(title, num_return_sequences=10, do_sample=True, top_k=0, top_p=0.5, temperature=1.8)
    
    # Extract just the translated text
    translations_text = [t['translation_text'] for t in translations]
    # Convert list of translations to a HTML string with rtl text direction
    translations_html = "<div dir='rtl'><ol>" + "".join(f"<li>{t}</li>" for t in translations_text) + "</ol></div>"
    return translations_html

# Define Gradio interface
iface = gr.Interface(
    fn=translate, 
    inputs=[gr.inputs.Textbox(lines=1, label="כותר", placeholder="Batman's Gun"), gr.Radio(["⚡ חיפוש אלומה", "💥 דגימה"], value="⚡ חיפוש אלומה", label="שיטת פיענוח", info="אם אתם רוצים להתפרע, תבחרו ב-💥 דגימה. מומלץ לנסות את שתי השיטות.")], 
    outputs=gr.outputs.HTML(label="Translated Titles"),
    title="🎬🔄💩 Lost In Translation",
    description=description_text,
    allow_flagging='never',
    examples=[["Batman's Gun", "דגימה"], ["Once Upon A Time in Hollywood", "חיפוש אלומה"]],
    article=explanation)

# Launch the interface
iface.launch()