minor bugs fixed

examples/code_editor_scripts.py

@@ -19,19 +19,28 @@ class CodeEditor:
         
 
     def load_layout(self):
+        # Initialize session state for problem description and other fields if not already set
+        if 'problem' not in st.session_state:
+            st.session_state.problem = self.data["default_problem_description"]
+        if 'loss_system_prompt' not in st.session_state:
+            st.session_state.loss_system_prompt = self.data["default_loss_system_prompt"]
+        if 'instruction' not in st.session_state:
+            st.session_state.instruction = self.data["instruction"]
+
         col1, col2 = st.columns([1, 1])
         with col1:
-            self.problem = st.text_area("Problem description:", self.data["default_problem_description"], height=300)
+            st.session_state.problem = st.text_area("Problem description:", st.session_state.problem, height=300)
         with col2:
-            self.loss_system_prompt = st.text_area("Loss system prompt:", self.data["default_loss_system_prompt"], height=150)
-            self.instruction = st.text_area("Instruction for formatted LLM call:", self.data["instruction"], height=100)
+            st.session_state.loss_system_prompt = st.text_area("Loss system prompt:", st.session_state.loss_system_prompt, height=150)
+            st.session_state.instruction = st.text_area("Instruction for formatted LLM call:", st.session_state.instruction, height=100)
 
-        if "code_content" not in st.session_state:
+        # Assume the code content also needs to be persistent
+        if 'code_content' not in st.session_state:
             st.session_state.code_content = self.data["default_initial_solution"]
 
         def update_code_content(value):
             st.session_state.code_content = value
-        
+            print(f"Code updated: {st.session_state.code_content}")
 
         col1, col2 = st.columns(2)
         with col1:
@@ -54,18 +63,6 @@ class CodeEditor:
                 )
 
 
-            # format_string = f"{instruction}\nProblem: {problem}\nCurrent Code: {st.session_state.code_content}"
-            # mui.Typography(format_string)
-            
-            # mui.Typography("Final Snippet vs. Current Solution:", sx={"fontSize": "20px", "fontWeight": "bold"})
-            # editor.MonacoDiff(
-            #     original=self.data["default_target_solution"],
-            #     modified=st.session_state.code_content,
-            #     height=300,
-            # )
-        
-
-
     def _run(self):
         # Code is the variable of interest we want to optimize -- so requires_grad=True
         solution = st.session_state.code_content
@@ -74,7 +71,7 @@ class CodeEditor:
                         role_description="code instance to optimize")
 
         # We are not interested in optimizing the problem -- so requires_grad=False
-        problem = tg.Variable(self.problem, 
+        problem = tg.Variable(st.session_state.problem, 
                             requires_grad=False, 
                             role_description="the coding problem")
 
@@ -82,13 +79,13 @@ class CodeEditor:
         optimizer = tg.TGD(parameters=[code])
 
 
-        instruction = self.instruction
+        instruction = st.session_state.instruction
         llm_engine = self.llm_engine
-        loss_system_prompt = self.loss_system_prompt
+        loss_system_prompt = st.session_state.loss_system_prompt
         loss_system_prompt = tg.Variable(loss_system_prompt, requires_grad=False, role_description="system prompt to the loss function")
 
         format_string = "{instruction}\nProblem: {{problem}}\nCurrent Code: {{code}}"
-        format_string = format_string.format(instruction=self.instruction)
+        format_string = format_string.format(instruction=st.session_state.instruction)
 
         fields = {"problem": None, "code": None}
         formatted_llm_call = tg.autograd.FormattedLLMCall(engine=self.llm_engine,