Update app.py

app.py

@@ -1,139 +1,126 @@
-import os
-import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from peft import PeftModel
-from threading import Thread
+import os
+from huggingface_hub import InferenceClient
+
+# Setup HF Token
+token_path = os.path.expanduser("~/.cache/huggingface/token")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if not HF_TOKEN and os.path.exists(token_path):
+    with open(token_path) as f:
+        HF_TOKEN = f.read().strip()
+
+# Model Config - Using the STABLE base model for reliable Cloud Inference
+MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct" 
+client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
 
-BASE_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
-ADAPTER_MODEL_ID = "vsple/LegalBuddy-Qwen-1.5B"
+# THE REAL SYSTEM PROMPT
+system_prompt = """You are LegalBuddy, a professional legal document drafting assistant for Indian law. 
+Your objective is to help users generate highly accurate, structured legal documents.
 
-print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(
-    BASE_MODEL_ID,
-    trust_remote_code=True
-)
+STRICT INSTRUCTIONS:
+1. INITIAL LANGUAGE: Always start in English. 
+2. DYNAMIC LANGUAGE: If the user speaks in Hindi/Hinglish, you MUST respond in the same. Otherwise, stick to English.
+3. INTERVIEW MODE: Ask structured questions ONE AT A TIME to collect missing info (Landlord, Tenant, Rent, etc.).
+4. DRAFTING: When ready, generate the full professional legal document structure with # Headers and clear clauses.
+"""
 
-print("Loading base model...")
-base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL_ID,
-    device_map="auto",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    trust_remote_code=True
-)
+custom_css = """
+body, .gradio-container { font-family: 'Inter', -apple-system, sans-serif !important; background-color: #f8fafc !important; }
+#header { padding: 30px; background: linear-gradient(135deg, #1e293b 0%, #0f172a 100%); border-radius: 12px; margin-bottom: 25px; box-shadow: 0 4px 6px -1px rgba(0,0,0,0.1); text-align: center; }
+#header h1 { margin: 0; font-size: 32px; font-weight: 800; color: #ffffff !important; letter-spacing: -0.5px; }
+#header p { margin: 8px 0 0 0; font-size: 16px; color: #cbd5e1 !important; font-weight: 400; }
+.chatbot-container { border-radius: 12px !important; box-shadow: 0 10px 25px -5px rgba(0, 0, 0, 0.1) !important; background: white !important; }
+.message-wrap { font-size: 16px !important; line-height: 1.6 !important; }
+"""
 
-print("Loading adapter...")
-model = PeftModel.from_pretrained(
-    base_model,
-    ADAPTER_MODEL_ID
-)
-model.eval()
+def setup_chat(user_text, history):
+    history.append((user_text, ""))
+    return gr.update(value="", interactive=False), history, gr.update(visible=False), gr.update(visible=True)
 
-def predict(message, history, system_prompt, max_tokens, temperature, top_p):
+def chat_logic(history, temp, top_p_val, max_tokens):
     messages = [{"role": "system", "content": system_prompt}]
+    for u_msg, a_reply in history[:-1]:
+        if u_msg: messages.append({"role": "user", "content": u_msg})
+        if a_reply: messages.append({"role": "assistant", "content": a_reply})
     
-    # Process history (which is now a list of dictionaries)
-    for msg in history:
-        messages.append(msg)
-    
-    messages.append({"role": "user", "content": message})
-    
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+    messages.append({"role": "user", "content": history[-1][0]})
     
-    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,
-        temperature=temperature,
-        top_p=top_p
-    )
-    
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    
-    partial = ""
-    for new_text in streamer:
-        partial += new_text
-        yield partial
+    partial_response = ""
+    try:
+        response_stream = client.chat_completion(
+            messages,
+            max_tokens=int(max_tokens),
+            stream=True,
+            temperature=float(temp),
+            top_p=float(top_p_val),
+        )
+        for chunk in response_stream:
+            if chunk.choices and chunk.choices[0].delta.content:
+                partial_response += chunk.choices[0].delta.content
+                yield partial_response
+    except Exception as e:
+        yield f"⚠️ Connection Issue: {str(e)}"
 
-def user_fn(user_message, history):
-    # Returns the updated history as a list of dictionaries
-    return "", history + [{"role": "user", "content": user_message}]
+def process_interaction(chat_history, temp, top_p_val, max_tokens):
+    user_input = chat_history[-1][0]
+    for partial_response in chat_logic(chat_history, temp, top_p_val, max_tokens):
+        chat_history[-1] = (user_input, partial_response)
+        yield chat_history
 
-def bot_fn(history, system_prompt, max_tokens, temperature, top_p):
-    user_message = history[-1]["content"]
-    # Initialize assistant message
-    history.append({"role": "assistant", "content": ""})
-    
-    for response in predict(
-        user_message,
-        history[:-1],
-        system_prompt,
-        max_tokens,
-        temperature,
-        top_p
-    ):
-        history[-1]["content"] = response
-        yield history
+def finalize_chat():
+    return gr.update(interactive=True), gr.update(visible=True), gr.update(visible=False)
 
-with gr.Blocks(title="LegalBuddy AI Draft Engine") as demo:
-    gr.Markdown("# ⚖️ LegalBuddy: The Digital Legal Chamber")
-    
-    chatbot = gr.Chatbot(height=600)
-    
-    msg = gr.Textbox(
-        placeholder="Type your legal query or draft request..."
-    )
+with gr.Blocks(theme=gr.themes.Default(primary_hue="slate", neutral_hue="slate"), css=custom_css, title="LegalBuddy Pro") as demo:
+    with gr.Column(elem_id="header"):
+        gr.Markdown("<h1>LegalBuddy Pro</h1>\n<p>Professional Legal Drafting Assistant</p>")
     
     with gr.Row():
-        submit = gr.Button("Send", variant="primary")
-        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_slider = gr.Slider(1, 2048, value=1024, label="Max Tokens")
-        temp_slider = gr.Slider(0.1, 1.0, value=0.1, step=0.1, label="Temperature")
-        top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+        with gr.Column(scale=12): # Full Width
+            chatbot = gr.Chatbot(
+                height=650, 
+                show_label=False, 
+                show_copy_button=True, 
+                bubble_full_width=True, 
+                avatar_images=(None, "⚖️"),
+                elem_classes="chatbot-container"
+            )
+            
+            with gr.Row():
+                user_msg = gr.Textbox(
+                    show_label=False, 
+                    placeholder="I need a Rent Agreement for Mumbai...", 
+                    scale=9, 
+                    container=False,
+                    autofocus=True
+                )
+                submit_btn = gr.Button("Draft ➤", variant="primary", scale=1)
+                stop_btn = gr.Button("Stop 🛑", variant="stop", scale=1, visible=False)
+
+            with gr.Accordion("Advanced Settings", open=False):
+                with gr.Row():
+                    temp_s = gr.Slider(0.01, 1.0, 0.05, step=0.01, label="Temperature")
+                    top_p_s = gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top P")
+                    max_toks = gr.Slider(500, 4096, 2048, step=100, label="Max Tokens")
 
-    submit.click(
-        user_fn,
-        [msg, chatbot],
-        [msg, chatbot],
-        queue=False
+    # Wire up interactions
+    submit_event = submit_btn.click(
+        fn=setup_chat, inputs=[user_msg, chatbot], outputs=[user_msg, chatbot, submit_btn, stop_btn]
     ).then(
-        bot_fn,
-        [chatbot, system_prompt, max_tokens_slider, temp_slider, top_p_slider],
-        chatbot
+        fn=process_interaction, inputs=[chatbot, temp_s, top_p_s, max_toks], outputs=[chatbot]
+    ).then(
+        fn=finalize_chat, outputs=[user_msg, submit_btn, stop_btn]
     )
-    
-    msg.submit(
-        user_fn,
-        [msg, chatbot],
-        [msg, chatbot],
-        queue=False
+
+    user_msg.submit(
+        fn=setup_chat, inputs=[user_msg, chatbot], outputs=[user_msg, chatbot, submit_btn, stop_btn]
     ).then(
-        bot_fn,
-        [chatbot, system_prompt, max_tokens_slider, temp_slider, top_p_slider],
-        chatbot
+        fn=process_interaction, inputs=[chatbot, temp_s, top_p_s, max_toks], outputs=[chatbot]
+    ).then(
+        fn=finalize_chat, outputs=[user_msg, submit_btn, stop_btn]
     )
-    
-    clear.click(lambda: [], None, chatbot, queue=False)
+
+    stop_btn.click(fn=None, cancels=[submit_event])
 
 if __name__ == "__main__":
-    demo.queue()
-    demo.launch()
+    print("🚀 Launching LegalBuddy Pro (Full-Screen Chat)...")
+    demo.queue().launch(share=True, server_port=7865)