Expanded Excel functionality to support multiple sheet translation

app.py

@@ -4,6 +4,8 @@ from pathlib import Path
 #from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
 from transformers import M2M100ForConditionalGeneration
 from tokenization_small100 import SMALL100Tokenizer
+import io 
+
 
 st.set_page_config(page_title="Translation Demo", page_icon=":milky_way:", layout="wide")
 
@@ -27,13 +29,67 @@ def get_translation(src_code, trg_code, src):
 
 def open_input(the_file):
     
+    sheets = []
+    
     if the_file.name.endswith('.tsv'):
         parsed = pd.read_csv(the_file, sep="\t")
+        
     elif the_file.name.endswith('.xlsx'):
-        parsed = pd.read_excel(the_file)
+        xlsx = pd.ExcelFile(the_file)
+        if len(xlsx.sheet_names) > 1:
+            sheets = [sheet for sheet in xlsx.sheet_names]
+            parsed = [pd.read_excel(xlsx, sheet) for sheet in sheets]
+        else:
+            parsed = pd.read_excel(the_file)
+
+    return parsed, sheets
+
+
+def translate_data(data, valid_languages_col, valid_languages_col_trg, col_for_translation, languages):
+    translated_data = []
+    new_df = data
+    for text in data[col_for_translation]:
+        if len(text) > 0 and col_src_lang in languages and col_trg_lang in languages:
+            with st.spinner("Translating..."):
+                try:
+                    target_text = get_translation(valid_languages_col, valid_languages_col_trg, 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["SMALL-100 translation"] = translated_data
+            
+    return new_df
 
-    return parsed
 
+def select_column(data, is_excel=False):
+
+    if is_excel:
+        columns = (col for col in data[0].columns)
+    else:
+        columns = (col for col in data.columns)
+    
+    src_col = st.selectbox(
+        'Select the column to translate (WARNING: You can only select a single column - please make sure all columns are named accordingly):',
+        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")
+    
+    return submitted_cols, src_col, col_src_lang, col_trg_lang
+    
 
 st.subheader("SMALL-100 Translator")
 
@@ -82,60 +138,55 @@ with st.form("my_form"):
             st.write("Please enter the source text, source language and target language.")
 
 
-st.subheader('Input Excel/TSV')
+st.subheader('Input XLSX/TSV')
 uploaded_file = st.file_uploader("Choose a file")
 done = False
 
-
 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)
-    data = open_input(uploaded_file)
-    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")
+    valid_col = (lang for lang in valid_languages)
+    valid_col_trg = (lang for lang in valid_languages)
+    data, sheets = open_input(uploaded_file)
+
+    if len(sheets) > 0:
+        translated_sheets = []
+        submitted_cols, src_col, valid_col, valid_col_trg = select_column(data, is_excel=True)
+
+        if submitted_cols:
+            for sheet in data:
+                translated_sheets.append(translate_data(data, valid_col, valid_col_trg, src_col, valid_languages))
+                
+        done = True
         
+    else:
+        submitted_cols, src_col, valid_col, valid_col_trg = select_column(data)
+        
+        st.subheader("DataFrame")
+        st.write(data)
+        st.write(data.describe())
+
         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
+            new_df = translate_data(data, valid_col, valid_col_trg, src_col, valid_languages)
+        
+        done = True    
 
 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')
+    
+    if len(sheets) > 0:
+        pass
+        buffer = io.BytesIO()
+        with pd.ExcelWriter(buffer) as writer:
+            for idx, sheet in enumerate(translated_sheets):
+                sheet.to_excel(writer, sheet_name=sheets[idx])
+                
+        st.download_button('Download XLSX', buffer, 'translated_file.xlsx', 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', key='download-xlsx')
+
+    
+    else:    
+        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: