Update app.py

app.py

@@ -12,23 +12,24 @@ ADAPTER_MODEL_ID = "vsple/LegalBuddy-Qwen-1.5B"
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
 
-# Load model (Fixed loading logic for HF Spaces)
-print("Loading model and adapter...")
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = AutoModelForCausalLM.from_pretrained(
+# Load base model
+print("Loading base model...")
+base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL_ID,
-    torch_dtype=torch.float32,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto",
     trust_remote_code=True
-).to(device)
+)
 
-model = PeftModel.from_pretrained(model, ADAPTER_MODEL_ID)
+# Load adapter
+print("Loading adapter...")
+model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL_ID)
 model = model.eval()
 
 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})
+    for msg in history:
+        messages.append(msg)
     
     messages.append({"role": "user", "content": message})
     
@@ -58,7 +59,7 @@ def predict(message, history, system_prompt, max_tokens, temperature, top_p):
         partial_text += new_text
         yield partial_text
 
-# Custom theme (Fixed for Gradio 4+)
+# Custom theme
 theme = gr.themes.Soft(
     primary_hue="slate",
     secondary_hue="blue",
@@ -71,10 +72,11 @@ theme = gr.themes.Soft(
 with gr.Blocks(theme=theme, title="LegalBuddy AI Draft Engine") as demo:
     with gr.Row():
         gr.Markdown("# ⚖️ LegalBuddy: The Digital Legal Chamber")
+        gr.Markdown("### 🚀 Experience the future of AI-driven Legal Drafting")
     
     with gr.Row():
         with gr.Column(scale=2):
-            chatbot = gr.Chatbot(height=600, show_label=False)
+            chatbot = gr.Chatbot(height=600, show_label=False, type="messages")
             msg = gr.Textbox(
                 placeholder="Type your legal query or draft request here...",
                 container=False,
@@ -86,11 +88,11 @@ with gr.Blocks(theme=theme, title="LegalBuddy AI Draft Engine") as demo:
                 
             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.",
+                    value="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.",
                     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")
+                temp = gr.Slider(minimum=0.1, maximum=1.0, value=0.1, step=0.1, label="Drafting Precision (Temperature)")
                 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):
@@ -98,18 +100,20 @@ with gr.Blocks(theme=theme, title="LegalBuddy AI Draft Engine") as demo:
             draft_viewer = gr.Markdown(
                 label="Generated Legal Document",
                 container=True,
+                line_breaks=True,
+                header_links=True,
                 value="*The legal draft will appear here as you interact with the AI assistant...*"
             )
 
+    def user_msg(user_message, history):
+        return "", history + [{"role": "user", "content": user_message}]
+
     def bot_msg(history, system_prompt, max_tokens, temperature, top_p):
-        user_message = history[-1][0]
-        history[-1][1] = ""
+        user_message = history[-1]["content"]
+        history.append({"role": "assistant", "content": ""})
         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]]
+            history[-1]["content"] += token
+            yield history, history[-1]["content"]
 
     submit_btn.click(
         user_msg, [msg, chatbot], [msg, chatbot]
@@ -126,4 +130,4 @@ with gr.Blocks(theme=theme, title="LegalBuddy AI Draft Engine") as demo:
     clear_btn.click(lambda: ([], "*The legal draft will appear here...*"), None, [chatbot, draft_viewer])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()