Update pages/Translate_Text.py

pages/Translate_Text.py

@@ -9,14 +9,14 @@ import pandas as pd
 
 
 
-def translate_text_to_text(text, target_lang):
-    prefix = f'translate to {target_lang}: '
+def translate_text_to_text(text, source_lang, target_lang):
+    prefix = f'translate {source_lang} to {target_lang}: '
     text = prefix + text
-    tokenizer = AutoTokenizer.from_pretrained('CohereForAI/aya-101')
-    inputs = tokenizer(text, return_tensors='pt')
-    model = AutoModelForSeq2SeqLM.from_pretrained('CohereForAI/aya-101')
-    outputs = model.generate(inputs, max_new_tokens=len(inputs.input_ids[0]) * 3, do_sample=False)
-    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    tokenizer = AutoTokenizer.from_pretrained('stevhliu/my_awesome_opus_books_model')
+    input_ids = tokenizer(text, return_tensors='pt').input_ids
+    model = AutoModelForSeq2SeqLM.from_pretrained('stevhliu/my_awesome_opus_books_model')
+    output_ids = model.generate(input_ids, max_new_tokens=len(input_ids[0]) * 3, do_sample=False, top_k=30, top_p=0.95)
+    translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     return translated_text
 
 
@@ -35,59 +35,40 @@ def main():
     with im3:
         pass
     
-    languages = ['Afrikaans', 'Amharic', 'Arabic', 'Azerbaijani', 'Belarusian',
-                'Bengali', 'Bulgarian', 'Catalan', 'Cebuano', 'Czech', 'Welsh',
-                'Danish', 'German', 'Greek', 'English', 'Esperanto', 'Estonian',
-                'Basque', 'Finnish', 'Filipino', 'French', 'Western Frisian',
-                'Scottish Gaelic', 'Irish', 'Galician', 'Gujarati', 'Haitian',
-                'Hausa', 'Hebrew', 'Hindi', 'Hungarian', 'Armenian', 'Igbo',
-                'Indonesian', 'Icelandic', 'Italian', 'Javanese', 'Japanese',
-                'Kannada', 'Georgian', 'Kazakh', 'Khmer', 'Kirghiz', 'Korean',
-                'Kurdish', 'Lao', 'Latvian', 'Latin', 'Lithuanian', 'Luxembourgish',
-                'Malayalam', 'Marathi', 'Macedonian', 'Malagasy', 'Maltese',
-                'Mongolian', 'Maori', 'Malay', 'Burmese', 'Nepali', 'Dutch',
-                'Norwegian', 'Pedi', 'Nyanja', 'Odia', 'Panjabi', 'Persian',
-                'Polish', 'Portuguese', 'Pushto', 'Romanian', 'Russian', 'Sinhala',
-                'Slovak', 'Slovenian', 'Samoan', 'Shona', 'Sindhi', 'Somali',
-                'Southern Sotho', 'Spanish', 'Albanian', 'Serbian', 'Sundanese',
-                'Swahili', 'Swedish', 'Tamil', 'Telugu', 'Tajik', 'Thai', 'Turkish',
-                'Twi', 'Ukrainian', 'Urdu', 'Uzbek', 'Vietnamese', 'Xhosa', 'Yiddish',
-                'Yoruba', 'Chinese','Zulu'
-                ]
-    # languages = ['English', 'French']
-    # source_lang = st.sidebar.selectbox('Source Language', languages)
+    languages = ['English', 'French']
+    source_lang = st.sidebar.selectbox('Source Language', languages)
     target_lang = st.sidebar.selectbox('Target Language', languages, index=1)
     text = st.text_area('Text Translator', placeholder='Enter your input text here ...', height=200, label_visibility='hidden')
 
     if st.button('translate it'):
         if text != '':
-            # if (source_lang == 'English' and target_lang == 'English') or (source_lang == 'French' and target_lang == 'French'):
-            #     st.error('Expected different values for source and target languages, but got the same values!')
+            if (source_lang == 'English' and target_lang == 'English') or (source_lang == 'French' and target_lang == 'French'):
+                st.error('Expected different values for source and target languages, but got the same values!')
 
-            # else:
-            with st.expander('Original Text'):
-                st.write(text)
-                add_one_item(text, 'Text Translator')
+            else:
+                with st.expander('Original Text'):
+                    st.write(text)
+                    add_one_item(text, 'Text Translator')
 
-            with st.expander('Translated Text'):
-                translated_text = translate_text_to_text(text, target_lang)
-                st.write(translated_text)
-            
-            col1, col2 = st.columns(2)
-            with col1:
-                with st.expander('Download Translated Text'):
-                    FileDownloader(translated_text, 'txt').download()
+                with st.expander('Translated Text'):
+                    translated_text = translate_text_to_text(text, source_lang, target_lang)
+                    st.write(translated_text)
+                
+                col1, col2 = st.columns(2)
+                with col1:
+                    with st.expander('Download Translated Text'):
+                        FileDownloader(translated_text, 'txt').download()
 
-            with col2:
-                with st.expander('Translated Text Validation'):
-                    bleu_score = validate_translation(text, translated_text)
-                    df = pd.DataFrame({
-                        'Brevity Penalty' : bleu_score.bp,
-                        'the length of the original text' : bleu_score.ref_len,
-                        'the length of the translated text' : bleu_score.sys_len,
-                        'Ratio' : bleu_score.ratio
-                    }, index=1)
-                    st.dataframe(df)
+                with col2:
+                    with st.expander('Translated Text Validation'):
+                        bleu_score = validate_translation(text, translated_text)
+                        df = pd.DataFrame({
+                            'Brevity Penalty' : bleu_score.bp,
+                            'the length of the original text' : bleu_score.ref_len,
+                            'the length of the translated text' : bleu_score.sys_len,
+                            'Ratio' : bleu_score.ratio
+                        }, index=1)
+                        st.dataframe(df)
 
         else:
             st.error('Please enter a non-empty text.')