add modifiers. change default LM. change UI

app.py

@@ -17,15 +17,18 @@ from consts import ALL_BLOCKS, ALL_BOWLS
 from md_logger import MarkdownLogger
 
 default_open_ai_key = os.getenv('OPEN_AI_SECRET')
+chain_of_thought_affix = ' with a step by step explanation'
+ask_for_clarification_affix = ' or ask for clarification if you feel unclear'
 
-class DemoRunner:
+
+class DemoRunner(processing_utils.TempFileManager):
 
     def __init__(self):
         self._cfg = OmegaConf.to_container(OmegaConf.load('cfg.yaml'), resolve=True)
         self._env = None
         self._model_name = ''
         self._md_logger = MarkdownLogger()
-        self._temp_file_manager = processing_utils.TempFileManager()
+        processing_utils.TempFileManager.__init__(self)
 
     def make_LMP(self, env):
         # LMP env wrapper
@@ -114,7 +117,7 @@ class DemoRunner:
         for message in self._md_logger.get_messages():
             history.append((None, message))
         if self._env.cache_video:
-            temp_name = self._temp_file_manager.make_temp_copy_if_needed(video_file_name)
+            temp_name = self.make_temp_copy_if_needed(video_file_name)
             history.append((None, (temp_name, )))
 
         return self._md_logger.get_log(), self._env.get_camera_image(), video_file_name, history
@@ -130,7 +133,7 @@ def setup(api_key, model_name, n_blocks, n_bowls):
 
     info, img = demo_runner.setup(api_key, model_name, n_blocks, n_bowls)
     welcome_message = 'How can I help you?'
-    return info, img, demo_runner, [(None, welcome_message)]
+    return info, img, demo_runner, [(None, welcome_message)], None
 
 
 def run(demo_runner, chat_history):
@@ -143,74 +146,93 @@ def submit_chat(chat_message, history):
     history += [[chat_message, None]]
     return '', history
 
-if __name__ == '__main__':
-    with open('README.md', 'r') as f:
-        for _ in range(12):
-            next(f)
-        readme_text = f.read()
-
-    with gr.Blocks() as demo:
-        state = gr.State(None)
-
-        gr.Markdown(readme_text)
-        gr.Markdown('# Interactive Demo')
-        with gr.Row():
-            with gr.Column():
-                with gr.Row():
-                    inp_api_key = gr.Textbox(value=default_open_ai_key,
-                                             label='OpenAI API Key (this is not stored anywhere)', lines=1)
-                    inp_model_name = gr.Dropdown(label='Model Name', choices=['code-davinci-002', 'text-davinci-002'], value='code-davinci-002')
-                with gr.Row():
-                    inp_n_blocks = gr.Slider(label='Number of Blocks', minimum=0, maximum=5, value=3, step=1)
-                    inp_n_bowls = gr.Slider(label='Number of Bowls', minimum=0, maximum=5, value=3, step=1)
-
-                btn_setup = gr.Button("Setup/Reset Simulation")
-                info_setup = gr.Markdown(label='Setup Info')
-            with gr.Column():
-                img_setup = gr.Image(label='Current Simulation State')
-
-        with gr.Row():
-            with gr.Column():
-                chat_box = gr.Chatbot()
-                inp_instruction = gr.Textbox(label='Instruction', lines=1)
-                examples = gr.Examples(
-                    [
-                        'stack two of the blocks',
-                        'what color is the rightmost block?',
-                        'arrange the blocks into figure 3',
-                        'put blocks into non-matching bowls'
-                    ],
-                    inp_instruction,
-                )
-                btn_run = gr.Button("Run (this may take 30+ seconds)")
-                info_run = gr.Markdown(label='Generated Code')
-            with gr.Column():
-                video_run = gr.Video(label='Video of Last Instruction')
-
-        btn_setup.click(
-            setup,
-            inputs=[inp_api_key, inp_model_name, inp_n_blocks, inp_n_bowls],
-            outputs=[info_setup, img_setup, state, chat_box],
-        )
-        btn_run.click(
-            submit_chat,
-            [inp_instruction, chat_box],
-            [inp_instruction, chat_box],
-        ).then(
-            run,
-            inputs=[state, chat_box],
-            outputs=[info_run, img_setup, video_run, chat_box, inp_instruction],
-        )
-        inp_instruction.submit(
-            submit_chat,
-            [inp_instruction, chat_box],
-            [inp_instruction, chat_box],
-        ).then(
-            run,
-            inputs=[state, chat_box],
-            outputs=[info_run, img_setup, video_run,
-                     chat_box, inp_instruction],
-        )
+def add_cot(chat_messsage):
+    return chat_messsage.strip() + chain_of_thought_affix
+
+def add_clarification(chat_message):
+    return chat_message.strip() + ask_for_clarification_affix
+
+
+with open('README.md', 'r') as f:
+    for _ in range(12):
+        next(f)
+    readme_text = f.read()
+
+with gr.Blocks() as demo:
+    state = gr.State(None)
+
+    gr.Markdown(readme_text)
+    gr.Markdown('# Interactive Demo')
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                inp_api_key = gr.Textbox(value=default_open_ai_key,
+                                            label='OpenAI API Key (this is not stored anywhere)', lines=1)
+                inp_model_name = gr.Dropdown(label='Model Name', choices=['code-davinci-002', 'text-davinci-003', 'text-davinci-002'], value='text-davinci-003')
+            with gr.Row():
+                inp_n_blocks = gr.Slider(label='Number of Blocks', minimum=0, maximum=5, value=3, step=1)
+                inp_n_bowls = gr.Slider(label='Number of Bowls', minimum=0, maximum=5, value=3, step=1)
+
+            btn_setup = gr.Button("Setup/Reset Simulation")
+            info_setup = gr.Markdown(label='Setup Info')
+
+    with gr.Row():
+        with gr.Column():
+            chat_box = gr.Chatbot()
+            inp_instruction = gr.Textbox(label='Instruction', lines=1)
+            examples = gr.Examples(
+                [
+                    'stack two of the blocks',
+                    'what color is the rightmost block?',
+                    'arrange the blocks into figure 3',
+                    'put blocks into non-matching bowls',
+                    'swap the positions of one block and another',
+                ],
+                inp_instruction,
+            )
+            btn_add_cot = gr.Button(f'+{chain_of_thought_affix} (chain-of-thought)')
+            btn_add_cla = gr.Button(
+                f'+{ask_for_clarification_affix} (conversation)')
+            btn_run = gr.Button("Run (this may take 30+ seconds)")
+            info_run = gr.Markdown(label='Generated Code')
+        with gr.Column():
+            img_setup = gr.Image(label='Current Simulation State')
+            video_run = gr.Video(label='Most Recent Manipulation')
+
+    btn_setup.click(
+        setup,
+        inputs=[inp_api_key, inp_model_name, inp_n_blocks, inp_n_bowls],
+        outputs=[info_setup, img_setup, state, chat_box, video_run],
+    )
+    btn_add_cot.click(
+        add_cot,
+        inp_instruction,
+        inp_instruction,
+    )
+    btn_add_cla.click(
+        add_clarification,
+        inp_instruction,
+        inp_instruction,
+    )
+    btn_run.click(
+        submit_chat,
+        [inp_instruction, chat_box],
+        [inp_instruction, chat_box],
+    ).then(
+        run,
+        inputs=[state, chat_box],
+        outputs=[info_run, img_setup, video_run, chat_box, inp_instruction],
+    )
+    inp_instruction.submit(
+        submit_chat,
+        [inp_instruction, chat_box],
+        [inp_instruction, chat_box],
+    ).then(
+        run,
+        inputs=[state, chat_box],
+        outputs=[info_run, img_setup, video_run, chat_box, inp_instruction],
+    )
 
+if __name__ == '__main__':
     demo.queue()
     demo.launch()