Update main.py

main.py

@@ -4,65 +4,59 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from peft import PeftModel
 
 # -----------------------------
-# PAGE CONFIG + CUSTOM CSS
+# PAGE CONFIG + STYLE
 # -----------------------------
 st.set_page_config(page_title="Kurdish Translator", layout="wide")
 
-# Dark/Light Mode
 THEME = st.session_state.get("theme", "light")
 
-if THEME == "dark":
-    bg_color = "#111111"
-    text_color = "#ffffff"
-else:
-    bg_color = "#ffffff"
-    text_color = "#000000"
+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>
-    body {{
-        background-color: {bg_color};
-        color: {text_color};
-    }}
-    
-    /* 3 bouncing dots loader */
-    .dots-loader {{
-        display: flex;
-        justify-content: center;
-        margin-top: 20px;
-    }}
-    .dots-loader div {{
-        width: 12px;
-        height: 12px;
-        margin: 4px;
-        background-color: #4A90E2;
-        border-radius: 50%;
-        animation: bounce 0.6s infinite alternate;
-    }}
-    .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: rgba(0,0,0,0.05);
-        font-size: 18px;
-        min-height: 120px;
-    }}
-
-    </style>
-    """,
+<style>
+body {{
+    background-color: {bg};
+    color: {fg};
+}}
+
+.dots-loader {{
+    display: flex;
+    justify-content: center;
+    margin-top: 10px;
+}}
+.dots-loader div {{
+    width: 12px;
+    height: 12px;
+    margin: 4px;
+    background-color: #4A90E2;
+    border-radius: 50%;
+    animation: bounce 0.6s infinite alternate;
+}}
+.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;
+}}
+</style>
+""",
     unsafe_allow_html=True
 )
 
 # -----------------------------
-# LOAD MODEL ONCE (CACHED)
+# LOAD MODEL WITH LOADER
 # -----------------------------
 BASE_MODEL = "facebook/nllb-200-distilled-600M"
 LORA_REPO = "junaid17/nllb-kurdish-lora"
@@ -70,9 +64,13 @@ LORA_REPO = "junaid17/nllb-kurdish-lora"
 
 @st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-    base = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
-    model = PeftModel.from_pretrained(base, LORA_REPO)
+    with st.spinner(""):
+        st.markdown("<div class='dots-loader'><div></div><div></div><div></div></div>", unsafe_allow_html=True)
+
+        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+        base = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
+        model = PeftModel.from_pretrained(base, LORA_REPO)
+
     return model, tokenizer
 
 
@@ -89,83 +87,72 @@ def translate(text, src, tgt):
         tokenizer=tokenizer,
         src_lang=src,
         tgt_lang=tgt,
-        max_length=256
+        max_length=256,
     )
-    out = translator(text)[0]["translation_text"]
-    return out
+    return translator(text)[0]["translation_text"]
 
 
 # -----------------------------
-# UI LOGIC
+# UI
 # -----------------------------
 st.title("Kurdish ↔ English Translator (NLLB + LoRA)")
 
-col1, col2, col3 = st.columns([1, 1, 1])
-
-# Language options
-LANG1 = st.session_state.get("LANG1", "English")
-LANG2 = st.session_state.get("LANG2", "Kurdish")
+col1, colSwap, colTheme, col2 = st.columns([1, 0.6, 0.8, 1])
 
 with col1:
-    st.subheader("From:")
-    src_lang = st.selectbox("", ["English", "Kurdish"], index=0 if LANG1 == "English" else 1)
+    src_lang = st.selectbox("From:", ["English", "Kurdish"])
+
+with colSwap:
+    swap = st.button("⇆ Swap", use_container_width=True)
+
+with colTheme:
+    toggle = st.button("🌙" if THEME == "light" else "☀️", use_container_width=True)
 
 with col2:
-    st.write("")
-    swap = st.button("⇆ Swap Languages", key="swap_lang")
+    tgt_lang = st.selectbox("To:", ["Kurdish", "English"])
 
-with col3:
-    st.subheader("To:")
-    tgt_lang = st.selectbox("", ["Kurdish", "English"], index=0 if LANG2 == "Kurdish" else 1)
 
-# Swap action
+# Handle swap
 if swap:
-    LANG1, LANG2 = LANG2, LANG1
-    st.session_state["LANG1"] = LANG1
-    st.session_state["LANG2"] = LANG2
+    src_lang, tgt_lang = tgt_lang, src_lang
+    st.session_state["LANG1"] = src_lang
+    st.session_state["LANG2"] = tgt_lang
     st.rerun()
 
-# Map language codes (NLLB)
-lang_map = {
+# Handle dark/light switch
+if toggle:
+    st.session_state["theme"] = "dark" if THEME == "light" else "light"
+    st.rerun()
+
+
+# Language mapping
+lang_codes = {
     "English": "eng_Latn",
     "Kurdish": "ckb_Arab"
 }
 
-src_code = lang_map[src_lang]
-tgt_code = lang_map[tgt_lang]
+src_code = lang_codes[src_lang]
+tgt_code = lang_codes[tgt_lang]
 
-# Input area
 text = st.text_area("Enter text:", height=180)
 
-# Translation button
-translate_btn = st.button("Translate")
-
 # -----------------------------
-# TRANSLATION OUTPUT
+# TRANSLATE BUTTON
 # -----------------------------
-if translate_btn:
+if st.button("Translate"):
     if not text.strip():
-        st.warning("Please enter some text...")
+        st.warning("Please enter some text.")
     else:
-        st.markdown("<div class='dots-loader'><div></div><div></div><div></div></div>", unsafe_allow_html=True)
+        # Show 3-dot loader ONLY during translation
+        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)
 
-        st.subheader("Output")
-
-        st.markdown(
-            f"<div class='output-box'>{output}</div>",
-            unsafe_allow_html=True
-        )
-
-        st.button("Copy Translation", on_click=lambda: st.session_state.update({"copied": output}))
+        # Remove loader
+        loader.empty()
 
-# -----------------------------
-# DARK / LIGHT MODE BUTTON
-# -----------------------------
-st.write("---")
-toggle = st.button("Toggle Dark / Light Mode")
+        st.subheader("Output")
+        st.markdown(f"<div class='output-box'>{output}</div>", unsafe_allow_html=True)
 
-if toggle:
-    st.session_state["theme"] = "dark" if THEME == "light" else "light"
-    st.rerun()
+        # Copy button
+        st.button("Copy Output", on_click=lambda: st.session_state.update({"copied": output}))