bug fixedd

examples/code_editor_scripts.py

@@ -7,6 +7,10 @@ import os
 class CodeEditor:
     def __init__(self, data) -> None:
         self.data = data
+        # Initialize only if not already set to ensure it retains the original content
+        if 'original_code_content' not in st.session_state:
+            st.session_state.original_code_content = self.data["default_initial_solution"]
+
         self.llm_engine = tg.get_engine("gpt-4o")
         print("="*50, "init", "="*50)
         self.loss_value = ""
@@ -39,29 +43,30 @@ class CodeEditor:
             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}")
+            if st.session_state.iteration == 0:
+                st.session_state.code_content = value
+            # print(f"Code updated: {st.session_state.code_content}")
 
         col1, col2 = st.columns(2)
         with col1:
-            with elements("monaco_editors_1"):
-                mui.Typography("Initial Solution:", sx={"fontSize": "20px", "fontWeight": "bold"})
-                editor.Monaco(
-                    height=300,
-                    defaultLanguage="python",
-                    defaultValue=st.session_state.code_content,
-                    onChange=update_code_content
-                )
-        with col2:
-            with elements("monaco_editors_2"):
-                mui.Typography("Current Solution:", sx={"fontSize": "20px", "fontWeight": "bold"})
-                editor.Monaco(
-                    height=300,
-                    defaultLanguage="python",
-                    value=st.session_state.code_content,
-                    options={"readOnly": True}  # Make the editor read-only
-                )
-
+            with elements("monaco_editors_widget_original"):
+                st.markdown(f"**Initial solution:**")
+                # code = editor.Monaco(
+                #     height=300,
+                #     defaultLanguage="python",
+                #     defaultValue=st.session_state.original_code_content,
+                #     onChange=update_code_content,
+                #     label="Initial Solution Viewer",
+                # )
+
+                code = st.text_area("Edit your code here:", st.session_state.original_code_content, height=300)
+                # Update session state when text changes
+                if code is not None and st.session_state.original_code_content != code:
+                    update_code_content(code)
+
+                # if st.session_state.code_content != code:
+                #     update_code_content(code)
+        # with col2:
 
     def _run(self):
         # Code is the variable of interest we want to optimize -- so requires_grad=True
@@ -114,21 +119,38 @@ class CodeEditor:
         st.session_state.results.append({
             'iteration': st.session_state.iteration,
             'loss_value': self.loss_value,
-            'gradients': self.gradients
+            'gradients': self.gradients,
+            'code_content': st.session_state.code_content,
         })
 
         tabs = st.tabs([f"Iteration {i+1}" for i in range(st.session_state.iteration)])
 
+        # Include Highlight.js library and a theme CSS
+        st.markdown("""
+        <link rel="stylesheet"
+            href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.5.0/styles/default.min.css">
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.5.0/highlight.min.js"></script>
+        <script>hljs.highlightAll();</script>
+        """, unsafe_allow_html=True)
+
         for i, tab in enumerate(tabs):
             with tab:
                 result = st.session_state.results[i]
                 st.markdown(f"Current iteration: **{result['iteration']}**")
+
+                st.markdown("### Current solution")
+                st.markdown(f"""
+                    <pre><code class="language-python">{result["code_content"]}</code></pre>
+                """, unsafe_allow_html=True)
+                
+
                 col1, col2 = st.columns([1, 1])
                 with col1:
-                    st.markdown("## Loss value")
+                    st.markdown("### Loss value")
+                    st.markdown("**Loss value is based on previous code.**")
                     st.markdown(result['loss_value'])
                 with col2:
-                    st.markdown("## Code gradients")
+                    st.markdown("### Code gradients")
                     for j, g in enumerate(result['gradients']):
-                        st.markdown(f"### Gradient {j}")
+                        # st.markdown(f"### Gradient {j}")
                         st.markdown(g.value)

examples/example_math_scripts.py

@@ -8,6 +8,15 @@ import os
 class MathSolution:
     def __init__(self, data) -> None:
         self.data = data
+        if 'default_initial_solution' not in st.session_state:
+            st.session_state.default_initial_solution = self.data["default_initial_solution"]
+        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"]
+        if 'current_solution' not in st.session_state:
+            st.session_state.current_solution = self.data["default_initial_solution"]
+
         self.llm_engine = tg.get_engine("gpt-4o")
         print("="*50, "init", "="*50)
         self.loss_value = ""
@@ -18,15 +27,22 @@ class MathSolution:
         tg.set_backward_engine(self.llm_engine, override=True)
 
     def load_layout(self):
+
+        def update_solution_content(value):
+            if st.session_state.iteration == 0:
+                st.session_state.current_solution = value
+            # print(f"Code updated: {st.session_state.code_content}")
+
         col1, col2 = st.columns([1, 1])
         with col1:
-            self.initial_solution = st.text_area("Initial solution:", self.data["default_initial_solution"], height=300)
-        with col2:
-            self.loss_system_prompt = st.text_area("Loss system prompt:", self.data["default_loss_system_prompt"], height=300)
-
-        if "current_solution" not in st.session_state:
-            st.session_state.current_solution = self.data["default_initial_solution"]
+            # self.initial_solution = st.text_area("Initial solution:", self.data["default_initial_solution"], height=300)
 
+            solution = st.text_area("Initial solution:", st.session_state.default_initial_solution, height=300)
+            # Update session state when text changes
+            if solution is not None and st.session_state.default_initial_solution != solution:
+                update_solution_content(solution)
+        with col2:
+            self.loss_system_prompt = st.text_area("Loss system prompt:", st.session_state.loss_system_prompt, height=300)
 
     def _run(self):
         # Set up the textgrad variables
@@ -71,10 +87,11 @@ class MathSolution:
                 st.markdown(result['response'])
                 col1, col2 = st.columns([1, 1])
                 with col1:
-                    st.markdown("## Loss value")
+                    st.markdown("### Loss value")
+                    st.markdown("**Loss value is based on previous code.**")
                     st.markdown(result['loss_value'])
                 with col2:
-                    st.markdown("## Code gradients")
+                    st.markdown("### Code gradients")
                     for j, g in enumerate(result['gradients']):
                         st.markdown(f"### Gradient")
                         st.markdown(g.value)