Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from peft import PeftModel
 from threading import Thread
+import time
 
 # Model configuration
 BASE_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
@@ -12,6 +13,7 @@ ADAPTER_MODEL_ID = "vsple/LegalBuddy-Qwen-1.5B"
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
 
 # Load base model
+print("Loading base model...")
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL_ID,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
@@ -20,28 +22,18 @@ base_model = AutoModelForCausalLM.from_pretrained(
 )
 
 # Load adapter
+print("Loading adapter...")
 model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL_ID)
 model = model.eval()
 
-def respond(
-    message,
-    history,
-    system_message="You are LegalBuddy, a professional legal assistant specializing in Indian Law and Document Drafting. Provide precise, legally compliant advice and draft clauses in a structured format.",
-    max_tokens=1024,
-    temperature=0.1,
-    top_p=0.9,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
+def predict(message, history, system_prompt, max_tokens, temperature, top_p):
+    messages = [{"role": "system", "content": system_prompt}]
+    for human, assistant in history:
+        messages.append({"role": "user", "content": human})
+        messages.append({"role": "assistant", "content": assistant})
+    
     messages.append({"role": "user", "content": message})
-
-    # Apply chat template for Qwen
+    
     prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
@@ -68,18 +60,72 @@ def respond(
         partial_text += new_text
         yield partial_text
 
-# Define the Gradio Interface
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are LegalBuddy, a professional legal assistant specializing in Indian Law and Document Drafting.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.1, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
-    ],
-    title="⚖️ LegalBuddy AI Draft Demo",
-    description="Live demo of LegalBuddy-Qwen-1.5B (Fine-tuned). Type your legal queries or drafting requests below."
+# Custom theme
+theme = gr.themes.Soft(
+    primary_hub_palette=gr.themes.colors.slate,
+    secondary_hub_palette=gr.themes.colors.blue,
+).set(
+    body_background_fill="*neutral_50",
+    block_background_fill="white",
+    block_border_width="1px",
 )
 
+with gr.Blocks(theme=theme, title="LegalBuddy AI Draft Engine") as demo:
+    with gr.Row():
+        gr.Markdown("# ⚖️ LegalBuddy: The Digital Legal Chamber")
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(height=600, show_label=False)
+            msg = gr.Textbox(
+                placeholder="Type your legal query or draft request here...",
+                container=False,
+                scale=7
+            )
+            with gr.Row():
+                submit_btn = gr.Button("Send Request", variant="primary")
+                clear_btn = gr.Button("Clear Session")
+                
+            with gr.Accordion("⚙️ Expert Settings", open=False):
+                system_msg = gr.Textbox(
+                    value="You are LegalBuddy, a professional legal assistant specializing in Indian Law and Document Drafting.",
+                    label="System Protocol"
+                )
+                max_tok = gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max Output Tokens")
+                temp = gr.Slider(minimum=0.1, maximum=1.0, value=0.1, step=0.1, label="Drafting Precision")
+                top_p_val = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p Sampling")
+
+        with gr.Column(scale=3):
+            gr.Markdown("## 📄 Live Draft Preview")
+            draft_viewer = gr.Markdown(
+                label="Generated Legal Document",
+                container=True,
+                value="*The legal draft will appear here as you interact with the AI assistant...*"
+            )
+
+    def bot_msg(history, system_prompt, max_tokens, temperature, top_p):
+        user_message = history[-1][0]
+        history[-1][1] = ""
+        for token in predict(user_message, history[:-1], system_prompt, max_tokens, temperature, top_p):
+            history[-1][1] += token
+            yield history, history[-1][1]
+
+    def user_msg(user_message, history):
+        return "", history + [[user_message, None]]
+
+    submit_btn.click(
+        user_msg, [msg, chatbot], [msg, chatbot]
+    ).then(
+        bot_msg, [chatbot, system_msg, max_tok, temp, top_p_val], [chatbot, draft_viewer]
+    )
+    
+    msg.submit(
+        user_msg, [msg, chatbot], [msg, chatbot]
+    ).then(
+        bot_msg, [chatbot, system_msg, max_tok, temp, top_p_val], [chatbot, draft_viewer]
+    )
+    
+    clear_btn.click(lambda: ([], "*The legal draft will appear here...*"), None, [chatbot, draft_viewer])
+
 if __name__ == "__main__":
     demo.launch()