Update app.py

app.py

@@ -1,133 +1,227 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from peft import PeftModel
 from threading import Thread
-import time
 
-# Model configuration
-BASE_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
-ADAPTER_MODEL_ID = "vsple/LegalBuddy-Qwen-1.5B"
+# ==================================
+# Model Configuration
+# ==================================
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
+BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
+ADAPTER_MODEL = "vsple/LegalBuddy-Qwen-1.5B"
 
-# Load model and adapter
-print("Loading model and adapter...")
-device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+
+print("Loading base model...")
 base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL_ID,
+    BASE_MODEL,
+    device_map="auto",
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    low_cpu_mem_usage=True,
     trust_remote_code=True
 )
 
-model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL_ID)
-model = model.to(device)
-model = model.eval()
+print("Loading adapter...")
+model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL)
+model.eval()
+
+# ==================================
+# Generation Function
+# ==================================
+
+def generate_stream(message, history, system_prompt, max_tokens, temperature, top_p):
 
-def predict(message, history, system_prompt, max_tokens, temperature, top_p):
     messages = [{"role": "system", "content": system_prompt}]
-    for msg in history:
-        messages.append(msg)
-    
+
+    for user, assistant in history:
+        messages.append({"role": "user", "content": user})
+        messages.append({"role": "assistant", "content": assistant})
+
     messages.append({"role": "user", "content": message})
-    
+
     prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
     )
-    
-    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    
-    generate_kwargs = dict(
-        inputs,
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
+
+    generation_kwargs = dict(
+        **inputs,
         streamer=streamer,
         max_new_tokens=max_tokens,
         do_sample=True,
-        top_p=top_p,
         temperature=temperature,
+        top_p=top_p
     )
-    
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
 
-    partial_text = ""
-    for new_text in streamer:
-        partial_text += new_text
-        yield partial_text
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    partial = ""
+
+    for token in streamer:
+        partial += token
+        yield partial
+
+
+# ==================================
+# Chat Functions
+# ==================================
+
+def add_user_message(message, history):
+    history.append([message, ""])
+    return "", history
+
+
+def generate_bot_message(history, system_prompt, max_tokens, temperature, top_p):
+
+    user_message = history[-1][0]
+
+    for token in generate_stream(
+        user_message,
+        history[:-1],
+        system_prompt,
+        max_tokens,
+        temperature,
+        top_p
+    ):
+        history[-1][1] = token
+        yield history, token
+
+
+def clear_chat():
+    return [], "*The legal draft will appear here...*"
+
+
+# ==================================
+# Theme
+# ==================================
 
-# Custom theme
 theme = gr.themes.Soft(
     primary_hue="slate",
     secondary_hue="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")
-        gr.Markdown("### 🚀 Experience the future of AI-driven Legal Drafting")
-    
+# ==================================
+# UI
+# ==================================
+
+with gr.Blocks(
+    title="LegalBuddy AI Draft Engine",
+    css="""
+    .draft-viewer {
+        max-height: 650px;
+        overflow-y: auto;
+        padding: 20px;
+    }
+    """
+) as demo:
+
+    gr.Markdown("# ⚖️ LegalBuddy AI")
+    gr.Markdown("### AI-Powered Legal Drafting Assistant")
+
     with gr.Row():
+
+        # Chat Section
         with gr.Column(scale=2):
-            chatbot = gr.Chatbot(height=600, show_label=False, type="messages")
+
+            chatbot = gr.Chatbot(
+                height=600,
+                bubble_full_width=False
+            )
+
             msg = gr.Textbox(
-                placeholder="Type your legal query or draft request here...",
-                container=False,
-                scale=7
+                placeholder="Type your legal drafting request...",
+                container=False
             )
+
             with gr.Row():
-                submit_btn = gr.Button("Send Request", variant="primary")
+                send_btn = gr.Button("Send", 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. Provide precise, legally compliant advice and draft clauses in a structured format.",
-                    label="System Protocol"
+
+                system_prompt = gr.Textbox(
+                    label="System Protocol",
+                    value="You are LegalBuddy, a professional legal assistant specializing in Indian law and legal document drafting."
+                )
+
+                max_tokens = gr.Slider(
+                    minimum=1,
+                    maximum=2048,
+                    value=1024,
+                    step=1,
+                    label="Max Tokens"
                 )
-                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 (Temperature)")
-                top_p_val = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p Sampling")
 
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.1,
+                    step=0.1,
+                    label="Temperature"
+                )
+
+                top_p = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.05,
+                    label="Top-p"
+                )
+
+        # Draft Viewer
         with gr.Column(scale=3):
-            gr.Markdown("## 📄 Live Draft Preview")
-            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}]
+            gr.Markdown("## 📄 Draft Preview")
 
-    def bot_msg(history, system_prompt, max_tokens, temperature, top_p):
-        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]["content"] += token
-            yield history, history[-1]["content"]
+            draft_output = gr.Markdown(
+                value="*Generated legal document will appear here...*",
+                elem_classes="draft-viewer"
+            )
 
-    submit_btn.click(
-        user_msg, [msg, chatbot], [msg, chatbot]
+    # ==================================
+    # Events
+    # ==================================
+
+    send_btn.click(
+        add_user_message,
+        [msg, chatbot],
+        [msg, chatbot]
     ).then(
-        bot_msg, [chatbot, system_msg, max_tok, temp, top_p_val], [chatbot, draft_viewer]
+        generate_bot_message,
+        [chatbot, system_prompt, max_tokens, temperature, top_p],
+        [chatbot, draft_output]
     )
-    
+
     msg.submit(
-        user_msg, [msg, chatbot], [msg, chatbot]
+        add_user_message,
+        [msg, chatbot],
+        [msg, chatbot]
     ).then(
-        bot_msg, [chatbot, system_msg, max_tok, temp, top_p_val], [chatbot, draft_viewer]
+        generate_bot_message,
+        [chatbot, system_prompt, max_tokens, temperature, top_p],
+        [chatbot, draft_output]
     )
-    
-    clear_btn.click(lambda: ([], "*The legal draft will appear here...*"), None, [chatbot, draft_viewer])
+
+    clear_btn.click(
+        clear_chat,
+        None,
+        [chatbot, draft_output]
+    )
+
+# ==================================
+# Launch
+# ==================================
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(theme=theme)