Update app.py

app.py

@@ -6,7 +6,7 @@ import pickle
 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 repo
+from ch09util import subsequent_mask, create_model  # Import from ch09util.py in the Space repo
 
 # Device setup
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -18,8 +18,8 @@ st.set_page_config(
     layout="centered"
 )
 
-# Model name
-MODEL_NAME = "amiguel/translateEn2FR"
+# Model repository name (corrected to the actual model repo)
+MODEL_NAME = "amiguel/custom-en2fr-transformer-v1"
 
 # Title with rocket emojis
 st.title("🚀 English to French Translator 🚀")
@@ -72,7 +72,7 @@ def load_model_and_resources(hf_token):
         
         login(token=hf_token)
         
-        # Load tokenizer
+        # Load tokenizer from the model repo
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_NAME,
             token=hf_token
@@ -108,7 +108,7 @@ def load_model_and_resources(hf_token):
             def forward(self, src, tgt, src_mask, tgt_mask, **kwargs):
                 return self.model(src, tgt, src_mask, tgt_mask)
 
-        # Load config with validation
+        # Load config with validation from the model repo
         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(
@@ -135,7 +135,7 @@ def load_model_and_resources(hf_token):
         
         model.eval()
 
-        # Load dictionaries from Hugging Face Hub
+        # Load dictionaries from the model repo
         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)