Update app.py

app.py

@@ -3,11 +3,10 @@ import torch
 import pandas as pd
 import PyPDF2
 import pickle
-import os
-from transformers import AutoTokenizer
-from huggingface_hub import login
+from transformers import AutoTokenizer, PreTrainedModel, PretrainedConfig
+from huggingface_hub import login, hf_hub_download
 import time
-from ch09util import subsequent_mask  # Ensure ch09util.py is available
+from utils.ch09util import subsequent_mask, create_model  # Ensure ch09util.py is available
 
 # Device setup
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -79,11 +78,10 @@ def load_model_and_resources(hf_token):
             token=hf_token
         )
         
-        # Load model
-        from transformers import PreTrainedModel, PretrainedConfig
+        # Define Transformer configuration
         class TransformerConfig(PretrainedConfig):
             model_type = "custom_transformer"
-            def __init__(self, src_vocab_size, tgt_vocab_size, d_model=256, d_ff=1024, h=8, N=6, dropout=0.1, **kwargs):
+            def __init__(self, src_vocab_size=11055, tgt_vocab_size=11239, d_model=256, d_ff=1024, h=8, N=6, dropout=0.1, **kwargs):
                 super().__init__(**kwargs)
                 self.src_vocab_size = src_vocab_size
                 self.tgt_vocab_size = tgt_vocab_size
@@ -93,11 +91,11 @@ def load_model_and_resources(hf_token):
                 self.N = N
                 self.dropout = dropout
 
+        # Define Transformer model
         class CustomTransformer(PreTrainedModel):
             config_class = TransformerConfig
             def __init__(self, config):
                 super().__init__(config)
-                from utils.ch09util import create_model
                 self.model = create_model(
                     config.src_vocab_size,
                     config.tgt_vocab_size,
@@ -110,18 +108,26 @@ 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)
 
-        config = TransformerConfig.from_pretrained(MODEL_NAME, token=hf_token)
+        # Load config with validation
+        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'. "
+                "Using defaults (11055, 11239). For accuracy, update the training script to save these values."
+            )
+            config = TransformerConfig()
+        else:
+            config = TransformerConfig(**config_dict)
+
+        # Load model
         model = CustomTransformer.from_pretrained(
             MODEL_NAME,
             config=config,
             token=hf_token
         ).to(DEVICE)
         
-        # Load dictionaries (assumes dict.p was uploaded to the model repo)
-        dict_path = "dict.p"
-        if not os.path.exists(dict_path):
-            st.error("Dictionary file (dict.p) not found. Please ensure it was uploaded to the model repository.")
-            return None
+        # Load dictionaries from Hugging Face Hub
+        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)