Update app.py

app.py

@@ -425,56 +425,59 @@ def chat_rag(
 
 
 # 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.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."
-            ),
-            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."
-            ),
-            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."
-            ),
-            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),
-                title="💬 RAG Chat",
-                description="Chat with the custom R1 model, enhanced with retrieval-augmented memory. The model retrieves relevant info for informed responses."
-            )
-        ]
-    )
+with gr.Blocks(css="""
+    body {background-color: #f5f5f5; font-family: Arial, sans-serif;}
+    .gradio-container {max-width: 1000px; margin: auto; background: white; padding: 20px; border-radius: 10px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);}
+    h1 {color: #333; text-align: center; font-size: 2rem;}
+    h2 {color: #444; margin-top: 10px; font-size: 1.5rem;}
+    .gr-tab {padding: 10px;}
+""") as demo:
+
+    gr.Markdown("# 🚀 QLoRA Fine-tuning & RAG Chat Demo")
+    gr.Markdown("Welcome to the enhanced **QLoRA fine-tuning and RAG-based chatbot interface**. This tool lets you fine-tune an AI model, generate text, and interact with a chatbot using retrieval-augmented responses.")
+
+    with gr.TabbedInterface():
+
+        # Fine-tuning tab
+        with gr.Tab(label="⚙️ Fine-tune Model"):
+            gr.Markdown("### Train your custom R1 model")
+            gr.Markdown("Fine-tune the model using QLoRA. This is **optional**, but recommended for better performance.")
+            finetune_btn = gr.Button("Start Fine-tuning")
+            finetune_output = gr.Textbox(label="Status", interactive=False)
+            finetune_btn.click(finetune_small_subset, inputs=None, outputs=finetune_output)
+
+        # Text Generation tab
+        with gr.Tab(label="✍️ Text Generation"):
+            gr.Markdown("### Generate text using your fine-tuned model")
+            input_prompt = gr.Textbox(label="Enter Prompt", placeholder="Type something here...", lines=3)
+            temp_slider = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")
+            topp_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+            min_tokens = gr.Slider(1, 1000, value=50, step=10, label="Min New Tokens")
+            max_tokens = gr.Slider(1, 1000, value=200, step=10, label="Max New Tokens")
+            generate_btn = gr.Button("Generate Text")
+            output_box = gr.Textbox(label="Generated Output", lines=8, interactive=False)
+            generate_btn.click(predict, inputs=[input_prompt, temp_slider, topp_slider, min_tokens, max_tokens], outputs=output_box)
+
+        # Model Comparison tab
+        with gr.Tab(label="🆚 Compare Models"):
+            gr.Markdown("### Compare text outputs from your fine-tuned model and the official model")
+            compare_prompt = gr.Textbox(label="Enter Comparison Prompt", placeholder="Enter a prompt here...", lines=3)
+            compare_temp = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")
+            compare_topp = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+            compare_min_tokens = gr.Slider(1, 1000, value=50, step=10, label="Min New Tokens")
+            compare_max_tokens = gr.Slider(1, 1000, value=200, step=10, label="Max New Tokens")
+            compare_btn = gr.Button("Compare Models")
+            compare_output1 = gr.Textbox(label="Custom Model Output", lines=6, interactive=False)
+            compare_output2 = gr.Textbox(label="Official Model Output", lines=6, interactive=False)
+            compare_btn.click(compare_models, inputs=[compare_prompt, compare_temp, compare_topp, compare_min_tokens, compare_max_tokens], outputs=[compare_output1, compare_output2])
+
+        # Chatbot tab
+        with gr.Tab(label="💬 AI Chatbot"):
+            gr.Markdown("### Chat with an AI assistant using retrieval-augmented generation (RAG)")
+            chatbot = gr.Chatbot(label="AI Chatbot", height=400)
+            chat_input = gr.Textbox(placeholder="Ask me anything...", lines=2)
+            chat_btn = gr.Button("Send")
+            chat_output = gr.Chatbot(label="Chat History")
+            chat_btn.click(chat_rag, inputs=[chat_input, chatbot, temp_slider, topp_slider, min_tokens, max_tokens], outputs=[chat_output, chatbot])
 
 demo.launch()