Update pages/Translate_Text.py

pages/Translate_Text.py

@@ -4,18 +4,26 @@ from configs.download_files import FileDownloader
 from configs.db_configs import add_one_item
 from streamlit.components.v1 import html
 from configs.html_features import set_image
+from sacrebleu.compat import corpus_bleu  
+import pandas as pd
 
-def translate_text_to_text(text, source_lang, target_lang):
-    prefix = f'translate {source_lang} to {target_lang}: '
+
+
+def translate_text_to_text(text, target_lang):
+    prefix = f'translate to {target_lang}: '
     text = prefix + text
-    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)
-    translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    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)
     return translated_text
 
 
+def validate_translation(original_text, translated_text):
+    return corpus_bleu(translated_text, [original_text])
+
+
 def main():
     st.title('Text Translator')
     im1, im2, im3 = st.columns([1, 5.3, 1])
@@ -26,26 +34,61 @@ def main():
         html(set_image(url), height=500, width=500)
     with im3:
         pass
-    languages = ['English', 'French']
-    source_lang = st.sidebar.selectbox('Source Language', languages)
+    
+    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)
     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!')
-
-            else:
-                with st.expander('Original Text'):
-                    st.write(text)
-                    add_one_item(text, 'Text Translator')
+            # 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!')
 
-                with st.expander('Translated Text'):
-                    translated_text = translate_text_to_text(text, source_lang, target_lang)
-                    st.write(translated_text)
+            # 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 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.')