Upload app.py

app.py

@@ -55,10 +55,16 @@ def load_model():
     """
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     
-    # Combine model parts and get the checkpoint path
-    checkpoint_path = combine_model_parts()
+    # Load model directly from checkpoint
+    checkpoint_path = "last.ckpt"  # Assuming the checkpoint is in the root directory
     
-    # Load the model from combined checkpoint using Lightning module
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(
+            f"Model checkpoint {checkpoint_path} not found. "
+            "Please ensure the model checkpoint file 'last.ckpt' is present in the root directory."
+        )
+    
+    # Load the model from checkpoint using Lightning module
     model = LitSmollmv2.load_from_checkpoint(
         checkpoint_path,
         model_config=SmollmConfig,