Update app.py

app.py

@@ -10,7 +10,8 @@ dialects = {"Palestinian/Jordanian": "P", "Syrian": "S", "Lebanese": "L", "Egypt
 # translator_en2ar = pipeline(task="translation", model="guymorlan/English2Dialect")
 translator_en2ar = MarianMTModel.from_pretrained("guymorlan/English2Dialect", output_attentions=True)
 tokenizer_en2ar = AutoTokenizer.from_pretrained("guymorlan/English2Dialect")
-translator_ar2en = pipeline(task="translation", model="guymorlan/Shami2English")
+translator_ar2en = MarianMTModel.from_pretrained("guymorlan/Shami2English", output_attentions=True)
+tokenizer_ar2en = AutoTokenizer.from_pretrained("guymorlan/Shami2English")
 transliterator = pipeline(task="translation", model="guymorlan/DialectTransliterator")
 
 speech_config = speechsdk.SpeechConfig(subscription=os.environ.get('SPEECH_KEY'), region=os.environ.get('SPEECH_REGION'))
@@ -28,7 +29,7 @@ def generate_diverging_colors(num_colors, palette='Set3'): # courtesy of ChatGPT
     return colors_hex
 
 
-def align_words(outputs, tokenizer, encoder_input_ids, decoder_input_ids, threshold=0.4):
+def align_words(outputs, tokenizer, encoder_input_ids, decoder_input_ids, threshold=0.4, skip_first_src=True):
     alignment = []
     for i, tok in enumerate(outputs.cross_attentions[2][0][7]):
         alignment.append([[i], (tok > threshold).nonzero().squeeze(-1).tolist()])
@@ -93,7 +94,7 @@ def align_words(outputs, tokenizer, encoder_input_ids, decoder_input_ids, thresh
 
     srchtml = []
     for i, token in enumerate(encoder_input_ids[0]):
-        if i == 0:
+        if skip_first_src and i == 0:
             continue
         if f"trg_{i}" in colordict:
             label = f"trg_{i}"
@@ -158,13 +159,42 @@ def translate_english(input_text, include):
 
     return palhtml, pal_out, sy_out, lb_out, eg_out
 
-def translate_arabic(input_text):
+def translate_arabic(input_text, include=["Colorize"]):
     if not input_text:
         return ""
 
-    result = translator_ar2en([input_text])
-    return result[0]["translation_text"]
+    input_tokens = tokenizer_ar2en(input_text, return_tensors="pt").input_ids
+    # print(input_tokens)
+    outputs =  translator_ar2en.generate(input_tokens)
+    # print(outputs)
+
+    encoder_input_ids = input_tokens[0].unsqueeze(0)
+    decoder_input_ids = outputs[0].unsqueeze(0)
 
+    decoded = tokenizer_en2ar.batch_decode(outputs, skip_special_tokens=True)
+    # print(decoded)
+
+    print(include)
+    if "Colorize" in include:
+        html_outputs = translator_ar2en(input_ids=encoder_input_ids, decoder_input_ids=decoder_input_ids)
+
+        # set dynamic threshold
+        # print(input_tokens, input_tokens.shape)
+        if input_tokens.shape[1] < 20:
+            threshold = 0.1
+        elif input_tokens.shape[1] < 30:
+            threshold = 0.01
+        else:
+            threshold = 0.05
+
+        print("threshold", threshold)
+
+        srchtml, tgthtml = align_words(html_outputs, tokenizer_ar2en, encoder_input_ids, decoder_input_ids, threshold, skip_first_src=False)
+        enhtml = f"<div style='direction: rtl'>{srchtml}</div><br><br><div>{tgthtml}</div>"
+    else:
+        enhtml = f"<div style='font-size: 30px;'>{decoded[0]}</div>"
+
+    return enhtml
 
 def get_audio(input_text):
     audio_config = speechsdk.audio.AudioOutputConfig(filename=f"{input_text}.wav")
@@ -244,6 +274,7 @@ with gr.Blocks(title = "Levantine Arabic Translator", css=css, theme="default")
         input_text.submit(translate_english, inputs=[input_text, include], outputs=[pal_html, pal, sy, lb, eg],scroll_to_output=True)
         pal.change(get_transliteration, inputs=[pal, include], outputs=[pal_translit]);
         include.change(toggle_visibility, inputs=[include], outputs=[pal_translit, sy, lb, eg])
+
     with gr.Tab('Ar > En'):
         with gr.Row():
             with gr.Column():
@@ -252,8 +283,12 @@ with gr.Blocks(title = "Levantine Arabic Translator", css=css, theme="default")
                 btn = gr.Button("Translate", label="Translate")
                 gr.Markdown("Built by [Guy Mor-Lan](mailto:guy.mor@mail.huji.ac.il).")
             with gr.Column():
-                eng = gr.Textbox(label="English", lines=1, elem_id="liter")
+                with gr.Box(label = "English"):
+                    gr.Markdown("English")
+                    with gr.Box():
+                        eng = gr.HTML("<br>", label="English", elem_id="main")
         btn.click(translate_arabic,inputs=input_text, outputs=[eng])
+
     with gr.Tab("Transliterate"):
         with gr.Row():
             with gr.Column():