Update my_model/utilities/st_utils.py

my_model/utilities/st_utils.py

@@ -105,6 +105,10 @@ class StateManager:
         df = pd.DataFrame(data)
         st.table(df)
 
+    def get_model(self):
+        """Retrieve the KBVQA model from the session state."""
+        return st.session_state.get('kbvqa', None)
+
     def is_model_loaded(self):
         return 'kbvqa' in st.session_state and st.session_state['kbvqa'] is not None