Update main.py

main.py

@@ -4,24 +4,32 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from peft import PeftModel
 
 # -----------------------------
-# PAGE CONFIG + STYLE
+# PAGE CONFIG
 # -----------------------------
 st.set_page_config(page_title="Kurdish Translator", layout="wide")
 
-THEME = st.session_state.get("theme", "light")
-
+# Initialize session state
+if "theme" not in st.session_state:
+    st.session_state.theme = "light"
+if "src_lang" not in st.session_state:
+    st.session_state.src_lang = "English"
+if "tgt_lang" not in st.session_state:
+    st.session_state.tgt_lang = "Kurdish"
+if "output_text" not in st.session_state:
+    st.session_state.output_text = ""
+
+# Theme colors
+THEME = st.session_state.theme
 bg = "#111111" if THEME == "dark" else "#ffffff"
 fg = "#ffffff" if THEME == "dark" else "#000000"
 card_bg = "rgba(255,255,255,0.07)" if THEME == "dark" else "rgba(0,0,0,0.05)"
 
 st.markdown(
-    f"""
-<style>
+    f"""<style>
 body {{
     background-color: {bg};
     color: {fg};
 }}
-
 .dots-loader {{
     display: flex;
     justify-content: center;
@@ -37,31 +45,28 @@ body {{
 }}
 .dots-loader div:nth-child(2) {{ animation-delay: 0.2s; }}
 .dots-loader div:nth-child(3) {{ animation-delay: 0.4s; }}
-
 @keyframes bounce {{
     from {{ transform: translateY(0); }}
     to {{ transform: translateY(-12px); }}
 }}
-
 .output-box {{
     padding: 20px;
     border-radius: 12px;
     background: {card_bg};
     font-size: 18px;
     min-height: 120px;
+    word-wrap: break-word;
 }}
-</style>
-""",
+</style>""",
     unsafe_allow_html=True
 )
 
 # -----------------------------
-# LOAD MODEL (NO LOADER HERE)
+# LOAD MODEL
 # -----------------------------
 BASE_MODEL = "facebook/nllb-200-distilled-600M"
 LORA_REPO = "junaid17/nllb-kurdish-lora"
 
-
 @st.cache_resource
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
@@ -69,10 +74,8 @@ def load_model():
     model = PeftModel.from_pretrained(base, LORA_REPO)
     return model, tokenizer
 
-
 model, tokenizer = load_model()
 
-
 # -----------------------------
 # TRANSLATION FUNCTION
 # -----------------------------
@@ -87,6 +90,21 @@ def translate(text, src, tgt):
     )
     return translator(text)[0]["translation_text"]
 
+# -----------------------------
+# CALLBACK FUNCTIONS
+# -----------------------------
+def swap_languages():
+    st.session_state.src_lang, st.session_state.tgt_lang = (
+        st.session_state.tgt_lang,
+        st.session_state.src_lang
+    )
+
+def toggle_theme():
+    st.session_state.theme = "dark" if st.session_state.theme == "light" else "light"
+
+def copy_to_clipboard():
+    if st.session_state.output_text:
+        st.toast("✓ Text copied to clipboard!", icon="✅")
 
 # -----------------------------
 # UI
@@ -96,55 +114,67 @@ st.title("Kurdish ↔ English Translator (NLLB + LoRA)")
 col1, colSwap, colTheme, col2 = st.columns([1, 0.6, 0.8, 1])
 
 with col1:
-    src_lang = st.selectbox("From:", ["English", "Kurdish"])
+    src_lang = st.selectbox(
+        "From:",
+        ["English", "Kurdish"],
+        key="src_lang"
+    )
 
 with colSwap:
-    swap = st.button("⇆ Swap", use_container_width=True)
+    st.button("⇆ Swap", use_container_width=True, on_click=swap_languages)
 
 with colTheme:
-    toggle = st.button("🌙" if THEME == "light" else "☀️", use_container_width=True)
+    theme_icon = "🌙" if THEME == "light" else "☀️"
+    st.button(theme_icon, use_container_width=True, on_click=toggle_theme)
 
 with col2:
-    tgt_lang = st.selectbox("To:", ["Kurdish", "English"])
-
-# Swap
-if swap:
-    src_lang, tgt_lang = tgt_lang, src_lang
-    st.rerun()
-
-# Toggle theme
-if toggle:
-    st.session_state["theme"] = "dark" if THEME == "light" else "light"
-    st.rerun()
+    tgt_lang = st.selectbox(
+        "To:",
+        ["Kurdish", "English"],
+        key="tgt_lang"
+    )
 
 lang_codes = {
     "English": "eng_Latn",
     "Kurdish": "ckb_Arab"
 }
 
-src_code = lang_codes[src_lang]
-tgt_code = lang_codes[tgt_lang]
+src_code = lang_codes[st.session_state.src_lang]
+tgt_code = lang_codes[st.session_state.tgt_lang]
 
 text = st.text_area("Enter text:", height=180)
 
-
 # -----------------------------
-# TRANSLATE BUTTON (with loader)
+# TRANSLATE BUTTON
 # -----------------------------
-if st.button("Translate"):
+if st.button("Translate", type="primary", use_container_width=True):
     if not text.strip():
         st.warning("Please enter some text.")
     else:
-        loader = st.markdown(
-            "<div class='dots-loader'><div></div><div></div><div></div></div>",
-            unsafe_allow_html=True
-        )
-
-        output = translate(text, src_code, tgt_code)
-
-        loader.empty()
-
-        st.subheader("Output")
-        st.markdown(f"<div class='output-box'>{output}</div>", unsafe_allow_html=True)
-
-        st.button("Copy Output", on_click=lambda: st.session_state.update({"copied": output}))
+        with st.spinner("Translating..."):
+            try:
+                output = translate(text, src_code, tgt_code)
+                st.session_state.output_text = output
+            except Exception as e:
+                st.error(f"Translation error: {str(e)}")
+                st.session_state.output_text = ""
+
+# Display output if available
+if st.session_state.output_text:
+    st.subheader("Output")
+    st.markdown(
+        f"<div class='output-box'>{st.session_state.output_text}</div>",
+        unsafe_allow_html=True
+    )
+    
+    # Copy button with actual clipboard functionality
+    col_copy, col_empty = st.columns([1, 3])
+    with col_copy:
+        if st.button("📋 Copy Output", use_container_width=True):
+            # Use Streamlit's native clipboard through components
+            st.write(f"""
+                <script>
+                navigator.clipboard.writeText(`{st.session_state.output_text}`);
+                </script>
+            """, unsafe_allow_html=True)
+            st.toast("✓ Copied to clipboard!", icon="✅")