Update app.py

app.py

@@ -1,42 +1,48 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from peft import PeftModel
 from threading import Thread
 
-# ==================================
-# Model Configuration
-# ==================================
-
-BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
-ADAPTER_MODEL = "vsple/LegalBuddy-Qwen-1.5B"
+BASE_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
+ADAPTER_MODEL_ID = "vsple/LegalBuddy-Qwen-1.5B"
 
 print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+
+tokenizer = AutoTokenizer.from_pretrained(
+    BASE_MODEL_ID,
+    trust_remote_code=True,
+    cache_dir="/tmp/huggingface"
+)
 
 print("Loading base model...")
+
 base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
+    BASE_MODEL_ID,
     device_map="auto",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    trust_remote_code=True
+    dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    trust_remote_code=True,
+    cache_dir="/tmp/huggingface"
 )
 
 print("Loading adapter...")
-model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL)
+
+model = PeftModel.from_pretrained(
+    base_model,
+    ADAPTER_MODEL_ID,
+    cache_dir="/tmp/huggingface"
+)
+
 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 user, assistant in history:
+    for user, bot in history:
         messages.append({"role": "user", "content": user})
-        messages.append({"role": "assistant", "content": assistant})
+        messages.append({"role": "assistant", "content": bot})
 
     messages.append({"role": "user", "content": message})
 
@@ -46,7 +52,7 @@ def generate_stream(message, history, system_prompt, max_tokens, temperature, to
         add_generation_prompt=True
     )
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
 
     streamer = TextIteratorStreamer(
         tokenizer,
@@ -55,7 +61,7 @@ def generate_stream(message, history, system_prompt, max_tokens, temperature, to
     )
 
     generation_kwargs = dict(
-        **inputs,
+        inputs,
         streamer=streamer,
         max_new_tokens=max_tokens,
         do_sample=True,
@@ -68,25 +74,20 @@ def generate_stream(message, history, system_prompt, max_tokens, temperature, to
 
     partial = ""
 
-    for token in streamer:
-        partial += token
+    for new_text in streamer:
+        partial += new_text
         yield partial
 
 
-# ==================================
-# Chat Functions
-# ==================================
+def user(user_message, history):
+    return "", history + [[user_message, None]]
 
-def add_user_message(message, history):
-    history.append([message, ""])
-    return "", history
 
-
-def generate_bot_message(history, system_prompt, max_tokens, temperature, top_p):
+def bot(history, system_prompt, max_tokens, temperature, top_p):
 
     user_message = history[-1][0]
 
-    for token in generate_stream(
+    for response in predict(
         user_message,
         history[:-1],
         system_prompt,
@@ -94,134 +95,72 @@ def generate_bot_message(history, system_prompt, max_tokens, temperature, top_p)
         temperature,
         top_p
     ):
-        history[-1][1] = token
-        yield history, token
-
+        history[-1][1] = response
+        yield history
 
-def clear_chat():
-    return [], "*The legal draft will appear here...*"
 
+with gr.Blocks(title="LegalBuddy AI Draft Engine") as demo:
 
-# ==================================
-# Theme
-# ==================================
+    gr.Markdown("# ⚖️ LegalBuddy: The Digital Legal Chamber")
 
-theme = gr.themes.Soft(
-    primary_hue="slate",
-    secondary_hue="blue",
-)
-
-# ==================================
-# UI
-# ==================================
-
-with gr.Blocks(
-    title="LegalBuddy AI Draft Engine",
-    css="""
-    .draft-viewer {
-        max-height: 650px;
-        overflow-y: auto;
-        padding: 20px;
-    }
-    """
-) as demo:
+    chatbot = gr.Chatbot(height=600)
 
-    gr.Markdown("# ⚖️ LegalBuddy AI")
-    gr.Markdown("### AI-Powered Legal Drafting Assistant")
+    msg = gr.Textbox(
+        placeholder="Type your legal query or draft request..."
+    )
 
     with gr.Row():
-
-        # Chat Section
-        with gr.Column(scale=2):
-
-            chatbot = gr.Chatbot(
-                height=600,
-                bubble_full_width=False
-            )
-
-            msg = gr.Textbox(
-                placeholder="Type your legal drafting request...",
-                container=False
-            )
-
-            with gr.Row():
-                send_btn = gr.Button("Send", variant="primary")
-                clear_btn = gr.Button("Clear Session")
-
-            with gr.Accordion("⚙️ Expert Settings", open=False):
-
-                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"
-                )
-
-                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("## 📄 Draft Preview")
-
-            draft_output = gr.Markdown(
-                value="*Generated legal document will appear here...*",
-                elem_classes="draft-viewer"
-            )
-
-    # ==================================
-    # Events
-    # ==================================
-
-    send_btn.click(
-        add_user_message,
+        submit = gr.Button("Send")
+        clear = gr.Button("Clear")
+
+    with gr.Accordion("Advanced Settings", open=False):
+
+        system_prompt = gr.Textbox(
+            value="You are a professional legal assistant specializing in Indian law and legal document drafting.",
+            label="System Prompt"
+        )
+
+        max_tokens = gr.Slider(1, 2048, value=512, label="Max Tokens")
+
+        temperature = gr.Slider(
+            0.1,
+            1.0,
+            value=0.2,
+            step=0.1,
+            label="Temperature"
+        )
+
+        top_p = gr.Slider(
+            0.1,
+            1.0,
+            value=0.9,
+            step=0.05,
+            label="Top-p"
+        )
+
+    submit.click(
+        user,
         [msg, chatbot],
-        [msg, chatbot]
+        [msg, chatbot],
+        queue=False
     ).then(
-        generate_bot_message,
+        bot,
         [chatbot, system_prompt, max_tokens, temperature, top_p],
-        [chatbot, draft_output]
+        chatbot
     )
 
     msg.submit(
-        add_user_message,
+        user,
+        [msg, chatbot],
         [msg, chatbot],
-        [msg, chatbot]
+        queue=False
     ).then(
-        generate_bot_message,
+        bot,
         [chatbot, system_prompt, max_tokens, temperature, top_p],
-        [chatbot, draft_output]
-    )
-
-    clear_btn.click(
-        clear_chat,
-        None,
-        [chatbot, draft_output]
+        chatbot
     )
 
-# ==================================
-# Launch
-# ==================================
+    clear.click(lambda: None, None, chatbot, queue=False)
 
-if __name__ == "__main__":
-    demo.launch(theme=theme)
+demo.queue()
+demo.launch()