Update app.py

app.py

@@ -19,28 +19,46 @@ def download_argos_model(from_code, to_code):
     )
     argostranslate.package.install_from_path(package_to_install.download())
 
-# Create the app layout
+# App layout
 st.header("Text Machine Translation")
 input_text = st.text_input("Enter text to translate:")
-# Create a list of options for the select box
+
+# Language options and mappings
 options = ["German", "Romanian", "English", "French", "Spanish", "Italian"]
-langs = {"English":"en", "Romanian":"ro", "German":"de", "French":"fr", "Spanish":"es", "Italian":"it"}
-models = ["Helsinki-NLP", "t5-base", "t5-small", "t5-large", "Unbabel/Tower-Plus-2B", "Unbabel/TowerInstruct-Mistral-7B-v0.2", 'Google', 'Argos']
+langs = {"English": "en", "Romanian": "ro", "German": "de", "French": "fr", "Spanish": "es", "Italian": "it"}
+models = ["Helsinki-NLP", "t5-base", "t5-small", "t5-large", "Unbabel/Tower-Plus-2B", 
+          "Unbabel/TowerInstruct-Mistral-7B-v0.2", "Google", "Argos"]
+
+# Initialize session state if not already set
+if "sselected_language" not in st.session_state:
+    st.session_state["sselected_language"] = options[0]
+if "tselected_language" not in st.session_state:
+    st.session_state["tselected_language"] = options[1]
+
+# Create columns for language selection
+scol, swapcol, tcol = st.columns([3, 1, 3])
 
-# Create two columns
-scol, tcol = st.columns(2)
-# Place select boxes in columns
 with scol:
-    sselected_language = st.selectbox("Source language:", options, index=0, placeholder="Select source language")
+    sselected_language = st.selectbox("Source language:", options, 
+                                      index=options.index(st.session_state["sselected_language"]))
+with swapcol:
+    if st.button("🔄 Swap"):
+        # Swap languages in session state
+        st.session_state["sselected_language"], st.session_state["tselected_language"] = \
+            st.session_state["tselected_language"], st.session_state["sselected_language"]
+        st.experimental_rerun()  # Rerun to update UI
 with tcol:
-    tselected_language = st.selectbox("Target language:", options, index=1, placeholder="Select target language")
-model_name = st.selectbox("Select a model:", models, index=0, placeholder="Select language model")
+    tselected_language = st.selectbox("Target language:", options, 
+                                      index=options.index(st.session_state["tselected_language"]))
+
+# Model selection
+model_name = st.selectbox("Select a model:", models, index=0)
 
-sl = langs[sselected_language]
-tl = langs[tselected_language]
+# Language codes
+sl = langs[st.session_state["sselected_language"]]
+tl = langs[st.session_state["tselected_language"]]
 
-st.session_state["sselected_language"] = sselected_language
-st.session_state["tselected_language"] = tselected_language
+# Store selections in session state
 st.session_state["model_name"] = model_name
 
 if model_name == 'Helsinki-NLP':