Update app.py

app.py

@@ -19,7 +19,7 @@ from transformers import (
     pipeline,
 )
 
-NUM_EXAMPLES_FOR_FINETUNING = 75  # Constant for the number of examples to use for finetuning
+NUM_EXAMPLES_FOR_FINETUNING = 50  # Constant for the number of examples to use for finetuning
 TEXT_PIPELINE = None  # Global to store the custom R1 text generation pipeline
 COMPARISON_PIPELINE = None  # Global to store the official R1 text generation pipeline
 
@@ -424,70 +424,72 @@ def chat_rag(
     return history, history
 
 
-# Build the Gradio interface.
+# Build the Gradio interface with tabs.
 with gr.Blocks() as demo:
     gr.Markdown("# QLoRA Fine-tuning & RAG-based Chat Demo using Custom R1 Model")
     gr.Markdown("---")
 
-    gr.Markdown("## ⚙️ Fine-tuning (Optional)")
-    gr.Markdown("This section allows you to fine-tune the custom R1 model on a small subset of the ServiceNow dataset. This step is optional but can potentially improve the model's performance on ServiceNow-related tasks. **Note:** This process may take up to 5 minutes.")
-    finetune_btn = gr.Button("🚀 Start Fine-tuning (QLoRA)")
-    status_box = gr.Textbox(label="Fine-tuning Status", interactive=False)
-    finetune_btn.click(fn=finetune_small_subset, outputs=status_box)
-    gr.Markdown("---")
-
-    gr.Markdown("## ✍️ Direct Generation (No Retrieval)")
-    gr.Markdown("Enter a prompt below to generate text directly using the custom R1 model. This is standard text generation without retrieval augmentation.")
-    prompt_in = gr.Textbox(lines=3, label="Input Prompt", placeholder="Enter your prompt here...")
-    temperature = gr.Slider(0.0, 1.5, step=0.1, value=0.7, label="Temperature (Creativity)")
-    top_p = gr.Slider(0.0, 1.0, step=0.05, value=0.9, label="Top-p (Sampling Nucleus)")
-    min_tokens = gr.Slider(1, 2500, value=50, step=10, label="Min New Tokens")
-    max_tokens = gr.Slider(1, 2500, value=200, step=50, label="Max New Tokens")
-    output_box = gr.Textbox(label="Custom R1 Output", lines=8, interactive=False)
-    gen_btn = gr.Button("✨ Generate Text")
-    gen_btn.click(
-        fn=predict,
-        inputs=[prompt_in, temperature, top_p, min_tokens, max_tokens],
-        outputs=output_box
-    )
-    gr.Markdown("---")
-
-    gr.Markdown("## 🆚 Compare Custom R1 vs Official R1")
-    gr.Markdown("Enter a prompt to compare the text generation of your fine-tuned custom R1 model with the official DeepSeek-R1-Distill-Llama-8B model.")
-    compare_prompt_in = gr.Textbox(lines=3, label="Comparison Prompt", placeholder="Enter prompt for comparison...")
-    compare_btn = gr.Button("⚖️ Compare Models")
-    out_custom = gr.Textbox(label="Custom R1 Output", lines=6, interactive=False)
-    out_official = gr.Textbox(label="Official R1 Output", lines=6, interactive=False)
-    compare_btn.click(
-        fn=compare_models,
-        inputs=[compare_prompt_in, temperature, top_p, min_tokens, max_tokens],
-        outputs=[out_custom, out_official]
-    )
-    gr.Markdown("---")
-
-    gr.Markdown("## 💬 Chat with Retrieval-Augmented Memory (RAG)")
-    gr.Markdown("Chat with the custom R1 model, enhanced with a retrieval-augmented memory. The model will retrieve relevant information based on your queries to provide more informed responses.")
-    with gr.Row():
-        with gr.Column():
-            chatbot = gr.Chatbot(label="RAG Chatbot")
-            chat_state = gr.State([])
-            user_input = gr.Textbox(
-                show_label=False,
-                placeholder="Ask a question to the RAG Chatbot...",
-                lines=2
+    with gr.TabbedInterface(
+        [
+            gr.Interface(
+                fn=finetune_small_subset,
+                inputs=None,
+                outputs=gr.Textbox(label="Fine-tuning Status", interactive=False),
+                title="⚙️ Fine-tuning (Optional)",
+                description="This section allows you to fine-tune the custom R1 model on a small subset of the ServiceNow dataset. This step is optional but can potentially improve the model's performance on ServiceNow-related tasks. **Note:** This process may take up to 5 minutes.",
+                submit_button_text="🚀 Start Fine-tuning (QLoRA)"
+            ),
+            gr.Interface(
+                fn=predict,
+                inputs=[
+                    gr.Textbox(lines=3, label="Input Prompt", placeholder="Enter your prompt here..."),
+                    gr.Slider(0.0, 1.5, step=0.1, value=0.7, label="Temperature (Creativity)"),
+                    gr.Slider(0.0, 1.0, step=0.05, value=0.9, label="Top-p (Sampling Nucleus)"),
+                    gr.Slider(1, 2500, value=50, step=10, label="Min New Tokens"),
+                    gr.Slider(1, 2500, value=200, step=50, label="Max New Tokens")
+                ],
+                outputs=gr.Textbox(label="Custom R1 Output", lines=8, interactive=False),
+                title="✍️ Direct Generation",
+                description="Enter a prompt to generate text directly using the custom R1 model. This is standard text generation without retrieval augmentation.",
+                submit_button_text="✨ Generate Text"
+            ),
+            gr.Interface(
+                fn=compare_models,
+                inputs=[
+                    gr.Textbox(lines=3, label="Comparison Prompt", placeholder="Enter prompt for comparison..."),
+                    gr.Slider(0.0, 1.5, step=0.1, value=0.7, label="Temperature"),
+                    gr.Slider(0.0, 1.0, step=0.05, value=0.9, label="Top-p"),
+                    gr.Slider(1, 2500, value=50, step=10, label="Min New Tokens"),
+                    gr.Slider(1, 2500, value=200, step=50, label="Max New Tokens")
+                ],
+                outputs=[
+                    gr.Textbox(label="Custom R1 Output", lines=6, interactive=False),
+                    gr.Textbox(label="Official R1 Output", lines=6, interactive=False)
+                ],
+                title="🆚 Model Comparison",
+                description="Enter a prompt to compare the text generation of your fine-tuned custom R1 model with the official DeepSeek-R1-Distill-Llama-8B model.",
+                submit_button_text="⚖️ Compare Models"
+            ),
+            gr.ChatInterface(
+                fn=chat_rag,
+                chatbot=gr.Chatbot(label="RAG Chatbot"),
+                textbox=gr.Textbox(placeholder="Ask a question to the RAG Chatbot...", lines=2, show_label=False),
+                inputs=[
+                    "textbox",
+                    gr.State([]), # chat_state
+                    gr.Slider(0.0, 1.5, step=0.1, value=0.7, label="Temperature"),
+                    gr.Slider(0.0, 1.0, step=0.05, value=0.9, label="Top-p"),
+                    gr.Slider(1, 2500, value=50, step=10, label="Min New Tokens"),
+                    gr.Slider(1, 2500, value=200, step=50, label="Max New Tokens")
+                ],
+                title="💬 RAG Chat",
+                description="Chat with the custom R1 model, enhanced with retrieval-augmented memory. The model retrieves relevant info for informed responses.",
+                submit_button="➡️ Send",
+                clear_btn=None # Optional: You can add a clear button if needed
             )
-            send_btn = gr.Button("➡️ Send")
-    user_input.submit(
-        fn=chat_rag,
-        inputs=[user_input, chat_state, temperature, top_p, min_tokens, max_tokens],
-        outputs=[chat_state, chatbot]
-    )
-    send_btn.click(
-        fn=chat_rag,
-        inputs=[user_input, chat_state, temperature, top_p, min_tokens, max_tokens],
-        outputs=[chat_state, chatbot]
-    )
-    gr.Markdown("---")
+        ]
+    ).render():
+        pass # No need for extra elements outside the tabs now
 
 
 demo.launch()