Update main.py

main.py

@@ -1,62 +1,171 @@
 import streamlit as st
+import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from peft import PeftModel
 
-st.set_page_config(page_title="Kurdish Translator", layout="centered")
-st.title("Kurdish ↔ English Translator (NLLB + LoRA)")
+# -----------------------------
+# PAGE CONFIG + CUSTOM CSS
+# -----------------------------
+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"
+
+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;
+    }}
 
-# HuggingFace model repos
+    </style>
+    """,
+    unsafe_allow_html=True
+)
+
+# -----------------------------
+# LOAD MODEL ONCE (CACHED)
+# -----------------------------
 BASE_MODEL = "facebook/nllb-200-distilled-600M"
-LORA_REPO = "junaid17/nllb-kurdish-lora"   # <-- use HF model repo, NOT local folder
+LORA_REPO = "junaid17/nllb-kurdish-lora"
 
 
 @st.cache_resource
-def load_translation_pipeline(direction):
-    # Load tokenizer
+def load_model():
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    base = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
+    model = PeftModel.from_pretrained(base, LORA_REPO)
+    return model, tokenizer
 
-    # Load base NLLB model
-    base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
 
-    # Load LoRA adapter from HuggingFace Hub
-    model = PeftModel.from_pretrained(base_model, LORA_REPO)
+model, tokenizer = load_model()
 
-    # Language codes for NLLB
-    if direction == "English → Kurdish":
-        src = "eng_Latn"
-        tgt = "ckb_Arab"
-    else:
-        src = "ckb_Arab"
-        tgt = "eng_Latn"
 
-    # Build translation pipeline
+# -----------------------------
+# TRANSLATION FUNCTION
+# -----------------------------
+def translate(text, src, tgt):
     translator = pipeline(
         "translation",
         model=model,
         tokenizer=tokenizer,
         src_lang=src,
         tgt_lang=tgt,
-        max_length=256,
+        max_length=256
     )
+    out = translator(text)[0]["translation_text"]
+    return out
 
-    return translator
 
+# -----------------------------
+# UI LOGIC
+# -----------------------------
+st.title("Kurdish ↔ English Translator (NLLB + LoRA)")
 
-# UI
-direction = st.selectbox(
-    "Select Translation Direction",
-    ["English → Kurdish", "Kurdish → English"]
-)
+col1, col2, col3 = st.columns([1, 1, 1])
+
+# Language options
+LANG1 = st.session_state.get("LANG1", "English")
+LANG2 = st.session_state.get("LANG2", "Kurdish")
 
-text = st.text_area("Enter your text:", height=180)
+with col1:
+    st.subheader("From:")
+    src_lang = st.selectbox("", ["English", "Kurdish"], index=0 if LANG1 == "English" else 1)
 
-if st.button("Translate"):
+with col2:
+    st.write("")
+    swap = st.button("⇆ Swap Languages", key="swap_lang")
+
+with col3:
+    st.subheader("To:")
+    tgt_lang = st.selectbox("", ["Kurdish", "English"], index=0 if LANG2 == "Kurdish" else 1)
+
+# Swap action
+if swap:
+    LANG1, LANG2 = LANG2, LANG1
+    st.session_state["LANG1"] = LANG1
+    st.session_state["LANG2"] = LANG2
+    st.rerun()
+
+# Map language codes (NLLB)
+lang_map = {
+    "English": "eng_Latn",
+    "Kurdish": "ckb_Arab"
+}
+
+src_code = lang_map[src_lang]
+tgt_code = lang_map[tgt_lang]
+
+# Input area
+text = st.text_area("Enter text:", height=180)
+
+# Translation button
+translate_btn = st.button("Translate")
+
+# -----------------------------
+# TRANSLATION OUTPUT
+# -----------------------------
+if translate_btn:
     if not text.strip():
-        st.warning("Please enter some text.")
+        st.warning("Please enter some text...")
     else:
-        with st.spinner("Translating..."):
-            translator = load_translation_pipeline(direction)
-            result = translator(text)[0]["translation_text"]
+        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}))
+
+# -----------------------------
+# DARK / LIGHT MODE BUTTON
+# -----------------------------
+st.write("---")
+toggle = st.button("Toggle Dark / Light Mode")
 
-        st.subheader("Translation:")
-        st.success(result)
+if toggle:
+    st.session_state["theme"] = "dark" if THEME == "light" else "light"
+    st.rerun()