adding translation task

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 # Custom imports 
 from multipage import MultiPage
-from pages import text_gen, fill_mask, home_page # import your pages here
+from pages import text_gen, fill_mask, home_page, translation # import your pages here
 
 # Create an instance of the app 
 app = MultiPage()
@@ -11,6 +11,7 @@ app = MultiPage()
 app.add_page("Home Page", home_page.app)
 app.add_page("Mask Fill", fill_mask.app)
 app.add_page("Text Generation", text_gen.app)
+app.add_page("Translation", translation.app)
 
 # app.add_page("Machine Learning", machine_learning.app)
 # app.add_page("Data Analysis",data_visualize.app)

pages/fill_mask.py

@@ -53,7 +53,7 @@ def app():
             unmasker = get_model(model)    
             st.success('Model loaded correctly!')
         
-            gen = st.info('Generating text...')
+            gen = st.info('Generating Mask...')
             answer = unmasker(prompt)     
             gen.empty()    
 

pages/translation.py

@@ -0,0 +1,52 @@
+import streamlit as st
+import time
+from multipage import MultiPage
+from transformers import pipeline
+import torch
+
+
+def app():
+    st.markdown('## Translation task')
+    st.write('Write something in English and AI will translate')
+    st.markdown('## ')
+
+    @st.cache(allow_output_mutation=True, suppress_st_warning =True, show_spinner=False)
+    def get_model(model):
+        return pipeline(task = 'translation', model = model)
+        
+    col1, col2 = st.columns([2,1])
+
+        
+    with col1:
+        prompt= st.text_area('Your prompt here',
+            '''What is the translation of this sentence?''') 
+            
+    with col2:
+        select_model = st.radio(
+            "Translate from English to:",
+            ('France', 'German', 'Cinese'), index = 0)
+        if select_model == 'France':
+            model = 'Helsinki-NLP/opus-mt-en-fr' 
+        elif select_model == 'German':
+            model = 'Helsinki-NLP/opus-mt-en-de' 
+        elif select_model == 'Spanish':
+            model = 'Helsinki-NLP/opus-mt-en-es' 
+
+        with st.spinner('Loading Model... (This may take a while)'):
+            generator = get_model(model)    
+            st.success('Model loaded correctly!')
+
+    with col1:    
+        gen = st.info('Generating text...')
+        answer = generator(prompt)                 
+        gen.empty()                      
+                        
+    lst = answer[0]['generated_text']
+    
+    t = st.empty()
+    for i in range(len(lst)):
+        t.markdown("#### %s..." % lst[0:i])
+        time.sleep(0.04)
+
+
+

requirements.txt

@@ -2,4 +2,5 @@ streamlit
 transformers
 altair
 pandas
+sacremoses
 torch