Update app.py

app.py

@@ -206,13 +206,22 @@ def search_similar_information_page(model):
                     st.write("No embedding could be generated for the search query.")
 
 # Page for translation
-def translate_text_page(tokenizer, model):
-    st.subheader("Translate from Shona to English")
-    text_to_translate = st.text_area("Enter Shona text to translate:")
-    if text_to_translate:
-        translated_text = translate_text(text_to_translate, tokenizer, model)
-        st.write("Translated text:")
-        st.write(translated_text)
+def translate_text_page(tokenizer_shona_to_english, model_shona_to_english, tokenizer_english_to_shona, model_english_to_shona):
+    st.subheader("Translate Text")
+    translation_direction = st.radio("Select translation direction", ("Shona to English", "English to Shona"))
+
+    if translation_direction == "Shona to English":
+        text_to_translate = st.text_area("Enter Shona text to translate:")
+        if text_to_translate:
+            translated_text = translate_text(text_to_translate, tokenizer_shona_to_english, model_shona_to_english)
+            st.write("Translated text:")
+            st.write(translated_text)
+    else:
+        text_to_translate = st.text_area("Enter English text to translate:")
+        if text_to_translate:
+            translated_text = translate_text(text_to_translate, tokenizer_english_to_shona, model_english_to_shona)
+            st.write("Translated text:")
+            st.write(translated_text)
 
 # Streamlit app
 def main():
@@ -222,12 +231,17 @@ def main():
     if 'fasttext_model' not in st.session_state:
         st.session_state['fasttext_model'] = load_fasttext_model("Fast_text_50_dim")
     
-    if 'translation_model' not in st.session_state:
-        st.session_state['tokenizer'], st.session_state['translation_model'] = load_translation_model("fine_tuned_shona_to_english_model")
+    if 'translation_model_shona_to_english' not in st.session_state:
+        st.session_state['tokenizer_shona_to_english'], st.session_state['translation_model_shona_to_english'] = load_translation_model("fine_tuned_shona_to_english_model")
+    
+    if 'translation_model_english_to_shona' not in st.session_state:
+        st.session_state['tokenizer_english_to_shona'], st.session_state['translation_model_english_to_shona'] = load_translation_model("english_shona_model")
     
     fasttext_model = st.session_state['fasttext_model']
-    tokenizer = st.session_state['tokenizer']
-    translation_model = st.session_state['translation_model']
+    tokenizer_shona_to_english = st.session_state['tokenizer_shona_to_english']
+    translation_model_shona_to_english = st.session_state['translation_model_shona_to_english']
+    tokenizer_english_to_shona = st.session_state['tokenizer_english_to_shona']
+    translation_model_english_to_shona = st.session_state['translation_model_english_to_shona']
 
     st.sidebar.title("Options")
     option = st.sidebar.radio("Select an option", ("Clean Dataset", "Train Word Embedding", "Generate Embeddings", "Translate Text"))
@@ -284,8 +298,8 @@ def main():
             search_similar_information_page(fasttext_model)
 
     elif option == "Translate Text":
-        st.header("Translate Text from Shona to English")
-        translate_text_page(tokenizer, translation_model)
+        st.header("Translate Text")
+        translate_text_page(tokenizer_shona_to_english, translation_model_shona_to_english, tokenizer_english_to_shona, translation_model_english_to_shona)
 
 if __name__ == "__main__":
     main()