Update app.py

app.py

@@ -50,8 +50,6 @@ def respond(
     max_tokens,
     temperature,
     top_p,
-    filename,
-    file_format
 ):
     messages = [{"role": "system", "content": system_message}]
 
@@ -77,35 +75,41 @@ def respond(
         response += token
         yield response
     
-    # Save the final response to the specified file format
+    return response, history + [(message, response)]
+
+# Function to handle file saving after generation
+def save_generated_file(response, filename, file_format):
     saved_file = save_file(response, filename, file_format)
-    yield response, history + [(message, response)], saved_file
-
-# Gradio interface
-demo = gr.ChatInterface(
-    fn=respond,
-    additional_inputs=[
-        gr.Textbox(value="", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-P",
-        ),
-        gr.Textbox(value="output", label="Filename"),
-        gr.Radio(["pdf", "docx", "txt"], label="File Format", value="pdf"),
-    ],
-    outputs=[
-        gr.Textbox(label="Generated Text"),
-        gr.State(value=[]),  # history
-        gr.File(label="Download File"),
-    ],
-    css=css,
-    theme="allenai/gradio-theme",
-)
+    return saved_file
+
+# Gradio interface using Blocks
+with gr.Blocks(css=css) as demo:
+    system_message = gr.Textbox(value="", label="System message")
+    max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
+    temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
+    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
+    filename = gr.Textbox(value="output", label="Filename")
+    file_format = gr.Radio(["pdf", "docx", "txt"], label="File Format", value="pdf")
+
+    message = gr.Textbox(label="User Message")
+    chat_history = gr.State(value=[])
+    response_output = gr.Textbox(label="Generated Response")
+    file_output = gr.File(label="Download File")
+
+    generate_button = gr.Button("Generate")
+    save_button = gr.Button("Save to File")
+
+    generate_button.click(
+        respond,
+        inputs=[message, chat_history, system_message, max_tokens, temperature, top_p],
+        outputs=[response_output, chat_history]
+    )
+
+    save_button.click(
+        save_generated_file,
+        inputs=[response_output, filename, file_format],
+        outputs=[file_output]
+    )
 
 if __name__ == "__main__":
     demo.launch()