Update app.py

app.py

@@ -1,58 +1,74 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel, PeftConfig
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import re
 
-# Load adapter config and base model
-adapter_id = "KomalNaseem/depression-chatbot-model"
-peft_config = PeftConfig.from_pretrained(adapter_id)
+model_path = "./depression_model_part1"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
+model.eval()
 
-# Load base model
-base_model = AutoModelForCausalLM.from_pretrained(
-    peft_config.base_model_name_or_path,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
+user_history = []
+turn_counter = 0
+MAX_TURNS_FOR_PREDICTION = 8
 
-# Load the LoRA adapter
-model = PeftModel.from_pretrained(base_model, adapter_id)
-model.eval()
+def chat(user_input):
+    global user_history, turn_counter
+    turn_counter += 1
+
+    user_history.append(f"Human: {user_input}")
+
+    last_turns = user_history[-4:]
+    prompt = "\n".join(last_turns) + "\nAI:"
+
+    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()
 
-# 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
+    user_history.append(f"AI: {response}")
+
+    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,
         )
-    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()
+        pred_text = tokenizer.decode(output_pred_ids[0], skip_special_tokens=True)
+
+        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 if depression_prob is not None else "Prediction after 8 turns"
+
+iface = gr.Interface(
+    fn=chat,
+    inputs=gr.Textbox(lines=2, label="Your Message"),
+    outputs=[gr.Textbox(label="AI Response"), gr.Textbox(label="Depression Probability")],
+    title="Depression Detection Chatbot",
+    description="Chat with the AI. After 8 turns it predicts depression probability."
+)
+
+if __name__ == "__main__":
+    iface.launch()