Update app.py

app.py

@@ -1,18 +1,27 @@
 import streamlit as st
 from transformers import pipeline
+import pandas as pd
 
-pipeline = pipeline(task="translation", model="guymorlan/English2Shami")
+dialects = {"Palestinian/Jordanian": "P", "Syrian": "S", "Lebanese": "L", "Egyptian": "E"}
+
+pipeline = pipeline(task="translation", model="guymorlan/English2Dialect")
 
 st.title("English to Levantine Arabic")
 
-num_translations = st.sidebar.selectbox("Number of Translations:", [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], index=2)
+num_translations = st.sidebar.selectbox("Number of Translations Per Dialect:", [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], index=0)
 input_text = st.text_input("Enter English text:")
 
 if input_text:
-    result = pipeline(input_text, max_length=1024, num_return_sequences=num_translations, num_beams=max(num_translations, 5))
-   # st.markdown("<div style='font-size:24px; text-align:right; direction:rtl;'>", unsafe_allow_html=True)
-    outs = list(set([x["translation_text"] for x in result]))
-    for i in range(len(outs)):
-        st.markdown(f"<div style='font-size:24px; text-align:right; direction:rtl;'>{outs[i]}</div>", unsafe_allow_html=True)
-        st.markdown("<br>", unsafe_allow_html=True)
-    #st.markdown("</div>", unsafe_allow_html=True)
+    inputs = [f"{val} {input_text}" for val in dialects.values()]
+    result = pipeline(inputs, max_length=1024, num_return_sequences=num_translations, num_beams=max(num_translations, 5))
+    #df = pd.DataFrame({"Dialect": [x for x in dialects.keys()],
+    #                   "Translation": [x["translation_text"] for x in result]})
+
+    for i in range(len(result)):
+        st.markdown(f"<div style='font-size:24px'><b>{list(dialects.keys())[i]}:</b></div>", unsafe_allow_html=True)
+        if num_translations > 1:
+            for j in range(num_translations):
+                st.markdown(f"<div style='font-size:24px; text-align:right; direction:rtl;'>{result[i][j]['translation_text']}</div>", unsafe_allow_html=True)
+        else:
+            st.markdown(f"<div style='font-size:24px; text-align:right; direction:rtl;'>{result[i]['translation_text']}</div>", unsafe_allow_html=True)
+        st.markdown("<br>", unsafe_allow_html=True)