Update AKSHAYRAJAA/inference.py

AKSHAYRAJAA/inference.py

@@ -25,27 +25,48 @@ TRANSFORMS = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
+# def load_model():
+#     """
+#     Loads the model with the vocabulary and checkpoint.
+#     """
+#     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
+
+#     st.write("Initializing the model...")
+#     model = get_model_instance(vocabulary)  # Initialize the model
+
+#     if can_load_checkpoint():
+#         st.write("Loading checkpoint...")
+#         load_checkpoint(model)
+#     else:
+#         st.write("No checkpoint found, starting with untrained model.")
+
+#     model.eval()  # Set the model to evaluation mode
+#     st.write("Model is ready for inference.")
+#     return model
 def load_model():
     """
     Loads the model with the vocabulary and checkpoint.
     """
-    st.write("Loading dataset and vocabulary...")
+    print("Loading dataset and vocabulary...")
     dataset = load_dataset()  # Load dataset to access vocabulary
     vocabulary = dataset.vocab  # Assuming 'vocab' is an attribute of the dataset
 
-    st.write("Initializing the model...")
+    print("Initializing the model...")
     model = get_model_instance(vocabulary)  # Initialize the model
 
     if can_load_checkpoint():
-        st.write("Loading checkpoint...")
+        print("Loading checkpoint...")
         load_checkpoint(model)
     else:
-        st.write("No checkpoint found, starting with untrained model.")
+        print("No checkpoint found, starting with untrained model.")
 
     model.eval()  # Set the model to evaluation mode
-    st.write("Model is ready for inference.")
+    print("Model is ready for inference.")
     return model
 
+
 def preprocess_image(image_path):
     """
     Preprocess the input image for the model.