add ui components for finetuning preview

nerdai · Sep 21, 2024 · f5fe43f · f5fe43f
1 parent dc5ceb6
commit f5fe43f
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 4 deletions.
diff --git a/arc_finetuning_st/streamlit/app.py b/arc_finetuning_st/streamlit/app.py
@@ -17,6 +17,8 @@ def startup() -> Tuple[Controller,]:
 (controller,) = startup()
 
 
+if "show_finetuning_preview_dialog" not in st.session_state:
+    st.session_state["show_finetuning_preview_dialog"] = True
 if "disable_continue_button" not in st.session_state:
     st.session_state["disable_continue_button"] = True
 if "disable_start_button" not in st.session_state:

diff --git a/arc_finetuning_st/streamlit/controller.py b/arc_finetuning_st/streamlit/controller.py
@@ -241,6 +241,7 @@ def _display_attempt(
 
     async def handle_finetuning_preview_click(self) -> None:
         if self._handler:
+            st.session_state.show_finetuning_preview_dialog = True
             ft_vars = await self._handler.ctx.get("prompt_vars")
 
             @st.dialog("Finetuning Example", width="large")
@@ -259,6 +260,26 @@ def _display_finetuning_example() -> None:
                         self._attempts[-1], len(self._attempts)
                     )
                 )
-                st.code(FINETUNING_DATASET_EXAMPLE_TEMPLATE.format(**ft_vars))
+                with st.container(height=500, border=False):
+                    save_col, close_col = st.columns([1, 1])
+                    with save_col:
+                        if st.button("Save", use_container_width=True):
+                            st.session_state.show_finetuning_preview_dialog = (
+                                False
+                            )
+                            st.rerun()
+                    with close_col:
+                        if st.button("Close", use_container_width=True):
+                            st.session_state.show_finetuning_preview_dialog = (
+                                False
+                            )
+                            st.rerun()
+
+                    st.code(
+                        FINETUNING_DATASET_EXAMPLE_TEMPLATE.format(**ft_vars),
+                        language="markdown",
+                        wrap_lines=True,
+                    )
 
-            _display_finetuning_example()
+            if st.session_state.show_finetuning_preview_dialog:
+                _display_finetuning_example()
diff --git a/arc_finetuning_st/workflows/prompts.py b/arc_finetuning_st/workflows/prompts.py
@@ -73,8 +73,7 @@
 The predicted output was found to be incorrect and a critique has been articulated offering a potential
 reason as to why it may have been a flawed prediction.
 
-Your task now to create a new prediction that corrects from the previous attempts. Use
-the last attempt and critique.
+Your task now to create a new prediction that corrects from the previous attempts. Use the last attempt and critique.
 
 EXAMPLES:
 {examples}