Update AKSHAYRAJAA/inference.py

AKSHAYRAJAA/inference.py

@@ -49,21 +49,32 @@ def load_model():
     """
     Loads the model with the vocabulary and checkpoint.
     """
-    print("Loading dataset and vocabulary...")
+    st.write("Loading dataset and vocabulary...")
     dataset = load_dataset()  # Load dataset to access vocabulary
     vocabulary = dataset.vocab  # Assuming 'vocab' is an attribute of the dataset
 
-    print("Initializing the model...")
+    st.write("Initializing the model...")
     model = get_model_instance(vocabulary)  # Initialize the model
 
     if can_load_checkpoint():
-        print("Loading checkpoint...")
-        load_checkpoint(model)
+        st.write("Loading checkpoint...")
+        checkpoint = torch.load(config.CHECKPOINT_FILE, map_location=DEVICE)
+        
+        # Print out the checkpoint layer sizes for debugging
+        print({k: v.shape for k, v in checkpoint['state_dict'].items()})
+
+        # Try loading the checkpoint with strict=False to ignore mismatched layers
+        try:
+            model.load_state_dict(checkpoint['state_dict'], strict=False)
+            st.write("Checkpoint loaded successfully.")
+        except RuntimeError as e:
+            st.write(f"Error loading checkpoint: {e}")
+            st.write("Starting with untrained model.")
     else:
-        print("No checkpoint found, starting with untrained model.")
+        st.write("No checkpoint found, starting with untrained model.")
 
     model.eval()  # Set the model to evaluation mode
-    print("Model is ready for inference.")
+    st.write("Model is ready for inference.")
     return model