Update app.py

app.py

@@ -17,25 +17,7 @@ with st.container():
     [Twitter](https://twitter.com/FrancescoDaimon)
     """)
 
-st.title('Upload your data')
-
-st.subheader('Input TSV/CSV')
-uploaded_file = st.file_uploader("Choose a file")
-with st.spinner("Loading..."):
-    if uploaded_file is not None:
-        if uploaded_file.name.endswith('.tsv'):
-            data = pd.read_csv(uploaded_file, sep="\t")
-        else:
-            data = pd.read_csv(uploaded_file)
-    
-        st.subheader("DataFrame")
-        st.write(data)
-        st.write(data.describe())
-    
-    else:
-        st.info("☝️ Upload a TSV/CSV file")
-    
-
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
 st.subheader("MBART-50 Translator")
 
 source = "In the beginning the Universe was created. This has made a lot of people very angry and been widely regarded as a bad move."
@@ -44,6 +26,7 @@ model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-m
 tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
 
 
+
 def get_translation(src_code, trg_code, src):
 
     tokenizer.src_lang = src_code
@@ -56,19 +39,29 @@ def get_translation(src_code, trg_code, src):
     
     return trg
 
-valid_languages = ['en_XX', 'fr_XX', 'de_DE', 'it_IT', 'es_XX']
+valid_languages = [
+    'ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX',
+    'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP',
+    'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN' 'zh_CN', 'af_ZA', 'az_AZ', 'bn_IN',
+    'fa_IR', 'he_IL', 'hr_HR', 'id_ID', 'ka_GE', 'km_KH', 'mk_MK', 'ml_IN', 'mn_MN', 'mr_IN',
+    'pl_PL', 'ps_AF', 'pt_XX', 'sv_SE', 'sw_KE', 'ta_IN', 'te_IN', 'th_TH', 'tl_XX', 'uk_UA',
+    'ur_PK', 'xh_ZA', 'gl_ES', 'sl_SI'
+    ]
+
+valid_languages_tuple = (lang for lang in valid_languages)
+valid_languages_tuple_trg = (lang for lang in valid_languages)
 
 with st.form("my_form"):
     left_c, right_c = st.columns(2)
-    with left_c:
-        src_lang = st.selectbox(
+    #with left_c:
+    src_lang = st.selectbox(
         'Source language',
-        ('en_XX', 'fr_XX', 'de_DE', 'it_IT', 'es_XX'),
+        valid_languages_tuple,
         )
-    with right_c:
-        trg_lang = st.selectbox(
+    #with right_c:
+    trg_lang = st.selectbox(
         'Target language',
-        ('fr_XX', 'en_XX', 'de_DE', 'it_IT', 'es_XX')
+        valid_languages_tuple_trg,
         )
     source = st.text_area("Source", value=source, height=130, placeholder="Enter the source text...")
 
@@ -88,29 +81,64 @@ with st.form("my_form"):
             st.write("Please enter the source text, source language and target language.")
 
 
-def local_css(file_name):
-    with open(file_name) as f:
-        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)
+st.subheader('Input TSV')
+uploaded_file = st.file_uploader("Choose a file")
+done = False
 
-local_css("style/style.css")
 
-# ---- CONTACT ----
-with st.container():
-    st.write("---")
-    st.header("Get in Touch With Me!")
-    st.write("##")
-
-    contact_form = """
-    <form action="https://formsubmit.co/[email protected]" method="POST">
-        <input type="hidden" name="_captcha" value="false">
-        <input type="text" name="name" placeholder="Your name" required>
-        <input type="email" name="email" placeholder="Your email" required>
-        <textarea name="message" placeholder="Your message here" required></textarea>
-        <button type="submit">Send</button>
-    </form>
-    """
-    left_column, right_column = st.columns(2)
-    with left_column:
-        st.markdown(contact_form, unsafe_allow_html=True)
-    with right_column:
-        st.empty()
+
+if uploaded_file is not None:
+    valid_languages_col = (lang for lang in valid_languages)
+    valid_languages_col_trg = (lang for lang in valid_languages)
+    if uploaded_file.name.endswith('.tsv'):
+        data = pd.read_csv(uploaded_file, sep="\t")
+        st.subheader("DataFrame")
+        st.write(data)
+        st.write(data.describe())
+        columns = (col for col in data.columns)
+        src_col = st.selectbox(
+        'Select the column to translate:',
+        columns,
+        )
+        if src_col:
+            col_src_lang = st.selectbox(
+        'Source language:',
+        valid_languages_col,
+        )
+            col_trg_lang = st.selectbox(
+        'Target language:',
+        valid_languages_col_trg,
+        )
+            submitted_cols = st.button("Translate column")
+            if submitted_cols:
+                translated_data = []
+                new_df = data
+                for text in data[src_col]:
+                    if len(text) > 0 and col_src_lang in valid_languages and col_trg_lang in valid_languages:
+                        with st.spinner("Translating..."):
+                            try:
+                                target_text = get_translation(col_src_lang, col_trg_lang, text)[0]
+                                translated_data.append(target_text)
+                            except:
+                                st.subheader("Translation failed :sad:")
+                                break
+                    
+                    else:
+                        st.write("Please enter the source text, source language and target language.")
+                
+                new_df[src_col] = translated_data
+                done = True
+        
+    else:
+        data = pd.read_csv(uploaded_file)
+
+if done:
+    st.subheader("Translated DataFrame")
+    st.write(new_df)
+    st.write(new_df.describe())
+    to_dl = new_df.to_csv(index=False, sep='\t').encode('utf-8')
+    st.download_button('Download TSV', to_dl, 'translated_file.tsv', 'text/tsv', key='download-tsv')
+
+    
+else:
+    st.info("☝️ Upload a TSV file")