Update main.py

main.py

@@ -1,7 +1,5 @@
 import streamlit as st
-import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from peft import PeftModel
+from backend import load_model, translate
 
 # -----------------------------
 # PAGE CONFIG
@@ -23,17 +21,35 @@ 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)"
+input_bg = "#1e1e1e" if THEME == "dark" else "#f5f5f5"
+button_bg = "#2d2d2d" if THEME == "dark" else "#e0e0e0"
 
 st.markdown(
     f"""<style>
-body {{
+[data-testid="stAppViewContainer"] {{
     background-color: {bg};
     color: {fg};
 }}
+[data-testid="stHeader"] {{
+    background-color: {bg};
+}}
+[data-testid="stToolbar"] {{
+    background-color: {bg};
+}}
+.stTextArea textarea {{
+    background-color: {input_bg} !important;
+    color: {fg} !important;
+}}
+.stSelectbox {{
+    color: {fg};
+}}
+h1, h2, h3, p, label {{
+    color: {fg} !important;
+}}
 .dots-loader {{
     display: flex;
     justify-content: center;
-    margin-top: 10px;
+    margin: 20px 0;
 }}
 .dots-loader div {{
     width: 12px;
@@ -56,6 +72,7 @@ body {{
     font-size: 18px;
     min-height: 120px;
     word-wrap: break-word;
+    color: {fg};
 }}
 </style>""",
     unsafe_allow_html=True
@@ -64,31 +81,11 @@ body {{
 # -----------------------------
 # 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)
-    base = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
-    model = PeftModel.from_pretrained(base, LORA_REPO)
-    return model, tokenizer
+def get_model():
+    return load_model()
 
-model, tokenizer = load_model()
-
-# -----------------------------
-# TRANSLATION FUNCTION
-# -----------------------------
-def translate(text, src, tgt):
-    translator = pipeline(
-        "translation",
-        model=model,
-        tokenizer=tokenizer,
-        src_lang=src,
-        tgt_lang=tgt,
-        max_length=256,
-    )
-    return translator(text)[0]["translation_text"]
+model = get_model()
 
 # -----------------------------
 # CALLBACK FUNCTIONS
@@ -102,10 +99,6 @@ def swap_languages():
 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
 # -----------------------------
@@ -151,13 +144,22 @@ if st.button("Translate", type="primary", use_container_width=True):
     if not text.strip():
         st.warning("Please enter some text.")
     else:
-        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 = ""
+        # Show loader
+        loader_placeholder = st.empty()
+        loader_placeholder.markdown(
+            "<div class='dots-loader'><div></div><div></div><div></div></div>",
+            unsafe_allow_html=True
+        )
+        
+        try:
+            output = translate(src_code, tgt_code, model, text)
+            st.session_state.output_text = output
+        except Exception as e:
+            st.error(f"Translation error: {str(e)}")
+            st.session_state.output_text = ""
+        finally:
+            # Remove loader
+            loader_placeholder.empty()
 
 # Display output if available
 if st.session_state.output_text:
@@ -167,9 +169,6 @@ if st.session_state.output_text:
         unsafe_allow_html=True
     )
     
-    # Copy button with pyperclip alternative
-    col_copy, col_empty = st.columns([1, 3])
-    with col_copy:
-        # Display the output in a code block that users can easily select and copy
-        st.code(st.session_state.output_text, language=None)
-        st.caption("👆 Select and copy the text above (Ctrl+C / Cmd+C)")
+    # Copy section
+    st.code(st.session_state.output_text, language=None)
+    st.caption("👆 Select and copy the text above")