Fix state model issue

apps/mlm.py

@@ -27,12 +27,13 @@ def app(state):
 
     # @st.cache(persist=False) # TODO: Make this work with mlm_state. Currently not supported.
     def predict(transformed_image, caption_inputs):
-        outputs = mlm_state.model(pixel_values=transformed_image, **caption_inputs)
-        indices = np.where(caption_inputs["input_ids"] == bert_tokenizer.mask_token_id)[
-            1
-        ][0]
+        outputs = mlm_state.mlm_model(pixel_values=transformed_image, **caption_inputs)
+        print(outputs.logits.shape)
+        indices = np.where(caption_inputs["input_ids"] == bert_tokenizer.mask_token_id)[1][0]
+        print(indices)
         preds = outputs.logits[0][indices]
         scores = np.array(preds)
+        print(scores)
         return scores
 
     # @st.cache(persist=False)
@@ -56,10 +57,10 @@ def app(state):
         image = plt.imread(image_path)
         mlm_state.mlm_image = image
 
-    if mlm_state.model is None:
+    if mlm_state.mlm_model is None:
         # Display Top-5 Predictions
         with st.spinner("Loading model..."):
-            mlm_state.model = load_model(mlm_checkpoints[0])
+            mlm_state.mlm_model = load_model(mlm_checkpoints[0])
 
     if st.button(
         "Get a random example",

apps/vqa.py

@@ -31,7 +31,7 @@ def app(state):
     # @st.cache(persist=False)
     def predict(transformed_image, question_inputs):
         return np.array(
-            vqa_state.model(pixel_values=transformed_image, **question_inputs)[0][0]
+            vqa_state.vqa_model(pixel_values=transformed_image, **question_inputs)[0][0]
         )
 
     # @st.cache(persist=False)
@@ -65,9 +65,9 @@ def app(state):
         image = plt.imread(image_path)
         vqa_state.vqa_image = image
 
-    if vqa_state.model is None:
+    if vqa_state.vqa_model is None:
         with st.spinner("Loading model..."):
-            vqa_state.model = load_model(vqa_checkpoints[0])
+            vqa_state.vqa_model = load_model(vqa_checkpoints[0])
     
     # Display Top-5 Predictions