Update app.py

app.py

@@ -1,75 +1,58 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig
+import gradio as gr
 
-model_id = "KomalNaseem/depression-chatbot-model"  # replace with your model path
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
-model.eval()
-
-user_history = []
-turn_counter = 0
-MAX_TURNS_FOR_PREDICTION = 8
+# Load adapter config and base model
+adapter_id = "KomalNaseem/depression-chatbot-model"
+peft_config = PeftConfig.from_pretrained(adapter_id)
 
-def chat(user_input):
-    global user_history, turn_counter
-    turn_counter += 1
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
 
-    user_history.append(f"Human: {user_input}")
-    last_turns = user_history[-4:]
-    prompt = "\n".join(last_turns) + "\nAI:"
+# Load base model
+base_model = AutoModelForCausalLM.from_pretrained(
+    peft_config.base_model_name_or_path,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    output_ids = model.generate(
-        **inputs,
-        max_new_tokens=50,
-        do_sample=False,
-        pad_token_id=tokenizer.eos_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-    )
-    response_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    response = response_text.split("AI:")[-1].strip()
-    user_history.append(f"AI: {response}")
+# Load the LoRA adapter
+model = PeftModel.from_pretrained(base_model, adapter_id)
+model.eval()
 
-    depression_prob = None
-    if turn_counter == MAX_TURNS_FOR_PREDICTION:
-        prediction_prompt = (
-            "\n".join(user_history[-8:]) +
-            "\nAI: Based on this conversation, what is the probability that the human has depression? "
-            "Please answer with a number between 0 and 1."
-        )
-        inputs_pred = tokenizer(prediction_prompt, return_tensors="pt").to(device)
-        output_pred_ids = model.generate(
-            **inputs_pred,
-            max_new_tokens=10,
-            do_sample=False,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
+# Chat function
+def chat_fn(user_input, history):
+    # Prepare input prompt by appending to history
+    history = history or []
+    prompt = ""
+    for turn in history:
+        prompt += f"User: {turn[0]}\nAI: {turn[1]}\n"
+    prompt += f"User: {user_input}\nAI:"
+
+    # Tokenize and generate response
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
-        pred_text = tokenizer.decode(output_pred_ids[0], skip_special_tokens=True)
-
-        import re
-        match = re.search(r"0?\.\d+", pred_text)
-        if match:
-            try:
-                depression_prob = float(match.group(0))
-            except:
-                depression_prob = None
-
-    return response, depression_prob
-
-chat_interface = gr.Interface(
-    fn=chat,
-    inputs=gr.Textbox(label="User Input"),
-    outputs=[
-        gr.Textbox(label="Bot Response"),
-        gr.Textbox(label="Depression Probability (after 8 turns)")
-    ],
-    title="Depression Detection Chatbot",
-    description="This chatbot detects depression by conversing for multiple turns. After 8 turns, it estimates the probability of depression."
-)
-
-chat_interface.launch()
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Extract AI response (cutting off prompt)
+    ai_response = decoded.split("User:")[-1].split("AI:")[-1].strip()
+    history.append((user_input, ai_response))
+    return history, history
+
+# Launch Gradio chat interface
+gr.ChatInterface(
+    fn=chat_fn,
+    title="🧠 Depression Detection Chatbot",
+    description="This chatbot engages in a conversation and detects signs of depression using LLaMA 3.2 with a fine-tuned LoRA adapter.",
+    theme="soft",
+    examples=["I'm feeling really low today.", "I can't sleep at night and feel hopeless."],
+).launch()