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  # 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"
@@ -119,13 +119,22 @@ def load_model_and_resources(hf_token):
         else:
             config = TransformerConfig(**config_dict)
 
-        # Load model
-        model = CustomTransformer.from_pretrained(
-            MODEL_NAME,
-            config=config,
-            token=hf_token
-        ).to(DEVICE)
+        # 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)
+        from safetensors.torch import load_file
+        state_dict = load_file(weights_path)
+        model.load_state_dict(state_dict)
+        
+        # Move model to the target device safely
+        if DEVICE == "cuda":
+            model = model.to_empty(device=DEVICE)  # Move structure to GPU
+            model.load_state_dict(state_dict)      # Reload weights on GPU
+        else:
+            model = model.to(DEVICE)  # CPU can handle direct move after loading weights
         
+        model.eval()
+
         # 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:
@@ -243,7 +252,7 @@ if prompt := st.chat_input("Enter text to translate into French..."):
                 
                 # Display metrics
                 st.caption(
-                    f"🔑 Input Tokens: {input_tokens} | Output Tokens: {output_tokens} | "
+                    f"🤖 Input Tokens: {input_tokens} | Output Tokens: {output_tokens} | "
                     f"🕒 Speed: {speed:.1f}t/s | 💰 Cost (USD): ${total_cost_usd:.4f} | "
                     f"💵 Cost (AOA): {total_cost_aoa:.4f}"
                 )