Update app.py

app.py

@@ -7,7 +7,7 @@ import pandas as pd
 import torch
 import time
 
-# Check if 'peft' is installed
+# Check if 'peft' is installed (though not used here, kept for potential future use)
 try:
     from peft import PeftModel, PeftConfig
 except ImportError:
@@ -23,16 +23,21 @@ st.set_page_config(
     layout="centered"
 )
 
-# Model names
-BASE_MODEL_NAME = "amiguel/en2fr-transformer"
-#MODEL_OPTIONS = {
-#    "Full Fine-Tuned": "amiguel/instruct_BERT-base-uncased_model", #"amiguel/playbook_FT",#"amiguel/SmolLM2-360M-concise-reasoning",
-#    "LoRA Adapter": "amiguel/SmolLM2-360M-concise-reasoning-lora",
-#    "QLoRA Adapter": "amiguel/SmolLM2-360M-concise-reasoning-qlora"  # Hypothetical, adjust if needed
-#}
+# Model name
+MODEL_NAME = "amiguel/en2fr-transformer"
+
+# Translation prompt template
+TRANSLATION_PROMPT = """
+You are a professional translator specializing in English-to-French translation. Translate the following text accurately and naturally into French, preserving the original meaning and tone:
+
+**Text to translate:**
+{input_text}
+
+**French translation:**
+"""
 
 # Title with rocket emojis
-st.title("🚀 Translator 🚀")
+st.title("🚀 English to French Translator 🚀")
 
 # Configure Avatars
 USER_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/9904d9a0d445ab0488cf7395cb863cce7621d897/USER_AVATAR.png"
@@ -44,13 +49,9 @@ with st.sidebar:
     hf_token = st.text_input("Hugging Face Token", type="password", 
                            help="Get your token from https://huggingface.co/settings/tokens")
     
-    st.header("Model Selection 🤖")
-    model_type = st.selectbox("Choose Model Type", list(MODEL_OPTIONS.keys()), index=0)
-    selected_model = MODEL_OPTIONS[model_type]
-    
     st.header("Upload Documents 📂")
     uploaded_file = st.file_uploader(
-        "Choose a PDF or XLSX file",
+        "Choose a PDF or XLSX file to translate",
         type=["pdf", "xlsx"],
         label_visibility="collapsed"
     )
@@ -78,7 +79,7 @@ def process_file(uploaded_file):
 
 # Model loading function
 @st.cache_resource
-def load_model(hf_token, model_type, selected_model):
+def load_model(hf_token):
     try:
         if not hf_token:
             st.error("🔐 Authentication required! Please provide a Hugging Face token.")
@@ -87,32 +88,15 @@ def load_model(hf_token, model_type, selected_model):
         login(token=hf_token)
         
         # Load tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME, token=hf_token)
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=hf_token)
         
-        # Load model based on type
-        if model_type == "Full Fine-Tuned":
-            # Load full fine-tuned model directly
-            model = AutoModelForCausalLM.from_pretrained(
-                selected_model,
-                torch_dtype=torch.bfloat16,
-                device_map="auto",
-                token=hf_token
-            )
-        else:
-            # Load base model and apply PEFT adapter
-            base_model = AutoModelForCausalLM.from_pretrained(
-                BASE_MODEL_NAME,
-                torch_dtype=torch.bfloat16,
-                device_map="auto",
-                token=hf_token
-            )
-            model = PeftModel.from_pretrained(
-                base_model,
-                selected_model,
-                torch_dtype=torch.bfloat16,
-                is_trainable=False,  # Inference mode
-                token=hf_token
-            )
+        # Load the full model (no adapters since we're using the base transformer)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            token=hf_token
+        )
         
         return model, tokenizer
         
@@ -121,8 +105,8 @@ def load_model(hf_token, model_type, selected_model):
         return None
 
 # Generation function with KV caching
-def generate_with_kv_cache(prompt, file_context, model, tokenizer, use_cache=True):
-    full_prompt = f"Analyze this context:\n{file_context}\n\nQuestion: {prompt}\nAnswer:"
+def generate_translation(input_text, model, tokenizer, use_cache=True):
+    full_prompt = TRANSLATION_PROMPT.format(input_text=input_text)
     
     streamer = TextIteratorStreamer(
         tokenizer, 
@@ -158,20 +142,19 @@ for message in st.session_state.messages:
             st.markdown(message["content"])
 
 # Chat input handling
-if prompt := st.chat_input("Ask your inspection question..."):
+if prompt := st.chat_input("Enter text to translate into French..."):
     if not hf_token:
         st.error("🔑 Authentication required!")
         st.stop()
 
-    # Load model if not already loaded or if model type changed
-    if "model" not in st.session_state or st.session_state.get("model_type") != model_type:
-        model_data = load_model(hf_token, model_type, selected_model)
+    # Load model if not already loaded
+    if "model" not in st.session_state:
+        model_data = load_model(hf_token)
         if model_data is None:
             st.error("Failed to load model. Please check your token and try again.")
             st.stop()
             
         st.session_state.model, st.session_state.tokenizer = model_data
-        st.session_state.model_type = model_type
     
     model = st.session_state.model
     tokenizer = st.session_state.tokenizer
@@ -181,27 +164,28 @@ if prompt := st.chat_input("Ask your inspection question..."):
         st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Process file
+    # Process file or use prompt directly
     file_context = process_file(uploaded_file)
+    input_text = file_context if file_context else prompt
     
-    # Generate response with KV caching
+    # Generate translation
     if model and tokenizer:
         try:
             with st.chat_message("assistant", avatar=BOT_AVATAR):
                 start_time = time.time()
-                streamer = generate_with_kv_cache(prompt, file_context, model, tokenizer, use_cache=True)
+                streamer = generate_translation(input_text, model, tokenizer, use_cache=True)
                 
                 response_container = st.empty()
                 full_response = ""
                 
                 for chunk in streamer:
-                    cleaned_chunk = chunk.replace("<think>", "").replace("</think>", "").strip()
+                    cleaned_chunk = chunk.strip()
                     full_response += cleaned_chunk + " "
                     response_container.markdown(full_response + "▌", unsafe_allow_html=True)
                 
                 # Calculate performance metrics
                 end_time = time.time()
-                input_tokens = len(tokenizer(prompt)["input_ids"])
+                input_tokens = len(tokenizer(input_text)["input_ids"])
                 output_tokens = len(tokenizer(full_response)["input_ids"])
                 speed = output_tokens / (end_time - start_time)
                 
@@ -222,6 +206,6 @@ if prompt := st.chat_input("Ask your inspection question..."):
                 st.session_state.messages.append({"role": "assistant", "content": full_response})
                 
         except Exception as e:
-            st.error(f"⚡ Generation error: {str(e)}")
+            st.error(f"⚡ Translation error: {str(e)}")
     else:
         st.error("🤖 Model not loaded!")