Update my_model/state_manager.py

my_model/state_manager.py

@@ -4,6 +4,9 @@ import streamlit as st
 from my_model.utilities.gen_utilities import free_gpu_resources
 from my_model.KBVQA import KBVQA, prepare_kbvqa_model
 
+# Create two columns with different widths
+col1, col2, col3 = st.columns([0.2, 0.6, 0.2])  
+
 
 class StateManager:
 
@@ -22,9 +25,6 @@ class StateManager:
 
 
     def set_up_widgets(self):
-        
-        # Create two columns with different widths
-        col1, col2, col3 = st.columns([0.2, 0.6, 0.2])  # Adjust the ratio as needed
 
         with col1:
             st.selectbox("Choose a method:", ["Fine-Tuned Model", "In-Context Learning (n-shots)"], index=0, key='method')
@@ -79,8 +79,8 @@ class StateManager:
             st.session_state['previous_state'] = {'method': st.session_state.method, 'detection_model': st.session_state.detection_model, 'confidence_level': st.session_state.confidence_level}
             
             st.session_state['button_label'] = "Reload Model"
-            st.text('button changed')
-            self.has_state_changed()
+            #st.text('button changed')
+            #self.has_state_changed()
             free_gpu_resources()
         except Exception as e:
             st.error(f"Error loading model: {e}")
@@ -109,7 +109,7 @@ class StateManager:
             if self.is_model_loaded():
                 prepare_kbvqa_model(only_reload_detection_model=True)
                 st.session_state['kbvqa'].detection_confidence = st.session_state.confidence_level
-                st.success("Model reloaded with updated settings and ready for inference.")
+                col1.success("Model reloaded with updated settings and ready for inference.")
             free_gpu_resources()
         except Exception as e:
             st.error(f"Error reloading detection model: {e}")