Update app.py

app.py

@@ -3,10 +3,11 @@ import torch
 import pandas as pd
 import PyPDF2
 import pickle
+import os
 from transformers import AutoTokenizer, PreTrainedModel, PretrainedConfig
 from huggingface_hub import login, hf_hub_download
 import time
-from ch09util import subsequent_mask, create_model  # Import from ch09util.py in the Space repo
+from ch09util import subsequent_mask, create_model
 
 # Device setup
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -18,9 +19,15 @@ st.set_page_config(
     layout="centered"
 )
 
-# Model repository name (corrected to the actual model repo)
+# Model repository name
 MODEL_NAME = "amiguel/custom-en2fr-transformer-v1"
 
+# Retrieve Hugging Face token from environment variable
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if not HF_TOKEN:
+    st.error("🔐 Hugging Face token not found in environment variables. Please set HF_TOKEN in Space secrets.")
+    st.stop()
+
 # Title with rocket emojis
 st.title("🚀 English to French Translator 🚀")
 
@@ -28,12 +35,8 @@ st.title("🚀 English to French Translator 🚀")
 USER_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/9904d9a0d445ab0488cf7395cb863cce7621d897/USER_AVATAR.png"
 BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/991f4c6e4e1dc7a8e24876ca5aae5228bcdb4dba/Ataliba_Avatar.jpg"
 
-# Sidebar configuration
+# Sidebar configuration (removed token input)
 with st.sidebar:
-    st.header("Authentication 🔒")
-    hf_token = st.text_input("Hugging Face Token", type="password", 
-                           help="Get your token from https://huggingface.co/settings/tokens")
-    
     st.header("Upload Documents 📂")
     uploaded_file = st.file_uploader(
         "Choose a PDF or XLSX file to translate",
@@ -64,18 +67,14 @@ def process_file(uploaded_file):
 
 # Custom model loading function
 @st.cache_resource
-def load_model_and_resources(hf_token):
+def load_model_and_resources():
     try:
-        if not hf_token:
-            st.error("🔐 Authentication required! Please provide a Hugging Face token.")
-            return None
-        
-        login(token=hf_token)
+        login(token=HF_TOKEN)
         
         # Load tokenizer from the model repo
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_NAME,
-            token=hf_token
+            token=HF_TOKEN
         )
         
         # Define Transformer configuration
@@ -109,7 +108,7 @@ def load_model_and_resources(hf_token):
                 return self.model(src, tgt, src_mask, tgt_mask)
 
         # Load config with validation from the model repo
-        config_dict = TransformerConfig.from_pretrained(MODEL_NAME, token=hf_token).to_dict()
+        config_dict = TransformerConfig.from_pretrained(MODEL_NAME, token=HF_TOKEN).to_dict()
         if "src_vocab_size" not in config_dict or "tgt_vocab_size" not in config_dict:
             st.warning(
                 f"Config at {MODEL_NAME}/config.json is missing 'src_vocab_size' or 'tgt_vocab_size'. "
@@ -121,7 +120,7 @@ def load_model_and_resources(hf_token):
 
         # Initialize model on meta device and load weights explicitly
         model = CustomTransformer(config)
-        weights_path = hf_hub_download(repo_id=MODEL_NAME, filename="model.safetensors", token=hf_token)
+        weights_path = hf_hub_download(repo_id=MODEL_NAME, filename="model.safetensors", token=HF_TOKEN)
         from safetensors.torch import load_file
         state_dict = load_file(weights_path)
         model.load_state_dict(state_dict)
@@ -136,7 +135,7 @@ def load_model_and_resources(hf_token):
         model.eval()
 
         # Load dictionaries from the model repo
-        dict_path = hf_hub_download(repo_id=MODEL_NAME, filename="dict.p", token=hf_token)
+        dict_path = hf_hub_download(repo_id=MODEL_NAME, filename="dict.p", token=HF_TOKEN)
         with open(dict_path, "rb") as fb:
             en_word_dict, en_idx_dict, fr_word_dict, fr_idx_dict = pickle.load(fb)
         
@@ -190,15 +189,11 @@ for message in st.session_state.messages:
 
 # Chat input handling
 if prompt := st.chat_input("Enter text to translate into French..."):
-    if not hf_token:
-        st.error("🔑 Authentication required!")
-        st.stop()
-
     # Load model and resources if not already loaded
     if "model" not in st.session_state:
-        model_data = load_model_and_resources(hf_token)
+        model_data = load_model_and_resources()
         if model_data is None:
-            st.error("Failed to load model. Please check your token and try again.")
+            st.error("Failed to load model. Please check the HF_TOKEN in Space secrets and try again.")
             st.stop()
             
         st.session_state.model, st.session_state.tokenizer, \