Update app.py

app.py

@@ -1 +1,86 @@
-import gradio as grimport torchfrom transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamerfrom peft import PeftModelfrom threading import Thread# Model configurationBASE_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"ADAPTER_MODEL_ID = "vsple/LegalBuddy-Qwen-1.5B"# Load tokenizertokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)# Load base modelbase_model = AutoModelForCausalLM.from_pretrained(    BASE_MODEL_ID,    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,    device_map="auto",    trust_remote_code=True)# Load adaptermodel = 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]})    messages.append({"role": "user", "content": message})    # Apply chat template for Qwen    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,        streamer=streamer,        max_new_tokens=max_tokens,        do_sample=True,        top_p=top_p,        temperature=temperature,    )        t = Thread(target=model.generate, kwargs=generate_kwargs)    t.start()    partial_text = ""    for new_text in streamer:        partial_text += new_text        yield partial_text# Define the Gradio Interfacedemo = 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.",    theme="soft")if __name__ == "__main__":    demo.launch()
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from peft import PeftModel
+from threading import Thread
+
+# Model configuration
+BASE_MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
+ADAPTER_MODEL_ID = "vsple/LegalBuddy-Qwen-1.5B"
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
+
+# Load base model
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL_ID,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+# Load 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]})
+
+    messages.append({"role": "user", "content": message})
+
+    # Apply chat template for Qwen
+    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,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        top_p=top_p,
+        temperature=temperature,
+    )
+    
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    partial_text = ""
+    for new_text in streamer:
+        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.",
+    theme="soft"
+)
+
+if __name__ == "__main__":
+    demo.launch()