Fix layout

app.py

@@ -65,8 +65,8 @@ following schema, including the leading and trailing "```json" and "```":
     }
 ]
 ```
-"""    
-    
+"""
+
     return format_instructions
 
 
@@ -208,7 +208,7 @@ def generate_cells(dataset_id):
     prompt = generate_eda_prompt(features, df, first_code)
     messages = [gr.ChatMessage(role="user", content=prompt)]
     yield messages + [gr.ChatMessage(role="assistant", content="⏳ _Starting task..._")]
-    
+
     prompt_messages = [{"role": "user", "content": prompt}]
     output = inference_client.chat_completion(
         messages=prompt_messages, stream=True, max_tokens=2500
@@ -230,7 +230,9 @@ def generate_cells(dataset_id):
     formatted_prompt = generate_mapping_prompt(GENERATED_TEXT)
     logging.info(formatted_prompt)
     prompt_messages = [{"role": "user", "content": formatted_prompt}]
-    yield messages + [gr.ChatMessage(role="assistant", content="⏳ _Generating notebook..._")]
+    yield messages + [
+        gr.ChatMessage(role="assistant", content="⏳ _Generating notebook..._")
+    ]
 
     output = inference_client.chat_completion(
         messages=prompt_messages, stream=False, max_tokens=2500
@@ -239,7 +241,6 @@ def generate_cells(dataset_id):
     logging.info("---> Model output")
     logging.info(cells_txt)
 
-
     commands = get_txt_from_output(cells_txt)
     html_code = f"<iframe src='https://huggingface.co/datasets/{dataset_id}/embed/viewer' width='80%' height='560px'></iframe>"
     # Adding dataset viewer on the first part
@@ -265,6 +266,7 @@ def generate_cells(dataset_id):
     )
     yield messages
 
+
 def write_notebook_file(dataset_id, history):
     if not GENERATED_TEXT:
         raise Exception("No generated notebook")
@@ -293,83 +295,85 @@ def write_notebook_file(dataset_id, history):
     return history
 
 
-with gr.Blocks() as demo:
+with gr.Blocks(fill_height=True) as demo:
     gr.Markdown("# 🤖 Dataset notebook creator 🕵️")
-    dataset_name = HuggingfaceHubSearch(
-        label="Hub Dataset ID",
-        placeholder="Search for dataset id on Huggingface",
-        search_type="dataset",
-        value="",
-    )
+    with gr.Row():
+        with gr.Column(scale=1):
+            dataset_name = HuggingfaceHubSearch(
+                label="Hub Dataset ID",
+                placeholder="Search for dataset id on Huggingface",
+                search_type="dataset",
+                value="",
+            )
 
-    @gr.render(inputs=dataset_name)
-    def embed(name):
-        if not name:
-            return gr.Markdown("### No dataset provided")
-        html_code = f"""
-        <iframe
-          src="https://huggingface.co/datasets/{name}/embed/viewer/default/train"
-          frameborder="0"
-          width="100%"
-          height="350px"
-        ></iframe>
-        """
-        return gr.HTML(value=html_code)
-
-    generate_cells_btn = gr.Button("Generate notebook")
-
-    chatbot = gr.Chatbot(
-        label="Results",
-        type="messages",
-        avatar_images=(
-            None,
-            None,
-        ),
-    )
+            @gr.render(inputs=dataset_name)
+            def embed(name):
+                if not name:
+                    return gr.Markdown("### No dataset provided")
+                html_code = f"""
+                <iframe
+                src="https://huggingface.co/datasets/{name}/embed/viewer/default/train"
+                frameborder="0"
+                width="100%"
+                height="350px"
+                ></iframe>
+                """
+                return gr.HTML(value=html_code)
+
+            with gr.Row():
+                generate_eda_btn = gr.Button("Generate EDA notebook")
+                generate_training_btn = gr.Button("Generate Training notebook")
+                generate_rag_btn = gr.Button("Generate RAG notebook")
+        with gr.Column():
+            chatbot = gr.Chatbot(
+                label="Results",
+                type="messages",
+                avatar_images=(
+                    None,
+                    None,
+                ),
+            )
 
-    generate_cells_btn.click(
+    generate_eda_btn.click(
         generate_cells,
         inputs=[dataset_name],
         outputs=[chatbot],
     )
 
-
-    with gr.Row(visible=False) as auth_page:
-        with gr.Column():
-            gr.Markdown(
-                "Want to push to hub? Enter your token ([settings](https://huggingface.co/settings/tokens)):"
-            )
-            token_box = gr.Textbox(
-                "", label="token", placeholder="hf_xxx", type="password"
-            )
-            auth_error = gr.Markdown("", visible=False)
-
-    push_btn = gr.Button("Push notebook to hub", visible=False)
-    output_lbl = gr.HTML(value="", visible=False)
-
-
-
-    def auth(token):
-        if not token:
-            return {
-                auth_error: gr.Markdown(value="", visible=False),
-                push_btn: gr.Button(visible=False),
-            }
-        return {
-            auth_error: gr.Markdown(value="", visible=False),
-            push_btn: gr.Button("Push notebook to hub", visible=True),
-        }
-
-    token_box.change(
-        auth,
-        inputs=token_box,
-        outputs=[auth_error, push_btn],
-    )
-
-    push_btn.click(
-        push_notebook,
-        inputs=[dataset_name, token_box],
-        outputs=output_lbl,
-    )
+    # with gr.Row(visible=False) as auth_page:
+    #     with gr.Column():
+    #         gr.Markdown(
+    #             "Want to push to hub? Enter your token ([settings](https://huggingface.co/settings/tokens)):"
+    #         )
+    #         token_box = gr.Textbox(
+    #             "", label="token", placeholder="hf_xxx", type="password"
+    #         )
+    #         auth_error = gr.Markdown("", visible=False)
+
+    # push_btn = gr.Button("Push notebook to hub", visible=False)
+    # output_lbl = gr.HTML(value="", visible=False)
+
+    # def auth(token):
+    #     if not token:
+    #         return {
+    #             auth_error: gr.Markdown(value="", visible=False),
+    #             push_btn: gr.Button(visible=False),
+    #         }
+    #     return {
+    #         auth_error: gr.Markdown(value="", visible=False),
+    #         push_btn: gr.Button("Push notebook to hub", visible=True),
+    #     }
+
+    # token_box.change(
+    #     auth,
+    #     inputs=token_box,
+    #     outputs=[auth_error, push_btn],
+    # )
+
+    # push_btn.click(
+    #     push_notebook,
+    #     inputs=[dataset_name, token_box],
+    #     outputs=output_lbl,
+    # )
 
 demo.launch()