Uploaded files from Cursor

AdviceGenerator.py

@@ -0,0 +1,68 @@
+from typing import List, Dict
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    AutoModelForCausalLM,
+    pipeline
+)
+
+class AdviceGenerator(object):
+    def __init__(self, llm):
+        self.llm = llm
+        self.role = {
+            "role": "system",
+            "content": (
+                "You are a supportive assistant (not a mental health professional). "
+                "Be concrete and tailor every response to the user's situation. "
+                "Requirements:\n"
+                "1) Begin with ONE empathetic sentence that mentions a key detail from the user's text (name, event, constraint).\n"
+                "2) Then give 3–5 numbered, practical tips. Each tip must reference the user's situation (use names/keywords when present).\n"
+                "3) If the user's text involves talking to someone (crush, friend, teacher, parent, boss), include a short **Script** block "
+                "   with two options (in-person and text), customized with any names from the user's text.\n"
+                "4) Add a **Try now (2 min)** micro-step.\n"
+                "5) End with ONE targeted follow-up question that references the user's situation.\n"
+                "Avoid platitudes and generic advice; avoid clinical instructions."
+            ),
+        }
+
+    def generate_advice(
+        self,
+        disorder: str,
+        user_text: str,
+        history: List[Dict[str, str]] = None,
+        max_history_msgs: int = 50,
+        max_tokens: int = 600,        # give enough headroom
+        temperature: float = 0.6,
+        top_p: float = 0.9,
+    ) -> Dict[str, str]:
+        
+        msgs = [self.role]
+
+        # preserve rolling chat history if available
+        if history:
+            msgs.extend(history[-max_history_msgs:])
+
+        # always append the new user input
+        msgs.append({
+            "role": "user",
+            "content": (
+                "Use the exact situation below to personalize your advice. "
+                "Extract the main goal or barrier from the text and ground each tip in it.\n\n"
+                f"Detected context: {disorder}\n"
+                f"User text: {user_text}\n\n"
+                "Follow the system instructions strictly. Do NOT ask vague questions first."
+            ),
+        })
+
+        try:
+            resp = self.llm.create_chat_completion(
+                messages=msgs,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False,
+            )
+            text = resp["choices"][0]["message"]["content"].strip()
+            return {"text": text}
+        except Exception as e:
+            return {"text": f"I'm here to listen. Could you tell me more about how \"{user_text}\" is affecting you?"}

IllnessClassifier.py

@@ -0,0 +1,19 @@
+from transformers import ( # pylint: disable=import-error
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    AutoModelForCausalLM,
+    pipeline
+)
+
+import logging
+
+class IllnessClassifier(object):
+    def __init__(self):
+        self.classifier = pipeline("text-classification", model="dsuram/distilbert-mentalhealth-classifier")
+
+    def forward(self, text: str):
+        output = self.classifier(text)[0]
+        disorder = output['label']
+        confidence = output['score']
+        logging.info(f"Disorder: {disorder}, Confidence: {confidence}")
+        return disorder, confidence

SIDetector.py

@@ -0,0 +1,19 @@
+from transformers import ( # pylint: disable=import-error
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    AutoModelForCausalLM,
+    pipeline
+)
+
+import logging
+
+class SIDetector(object):
+    def __init__(self):
+        self.classifier = pipeline("sentiment-analysis", model="sentinet/suicidality")
+
+    def forward(self, text: str):
+        output = self.classifier(text)[0]
+        suicidal = True if output['label'] == 'LABEL_1' else False
+        confidence = output['score']
+        logging.info(f"Suicidal: {suicidal}, Confidence: {confidence}")
+        return suicidal, confidence

UofTearsBot.py

@@ -0,0 +1,60 @@
+from AdviceGenerator import AdviceGenerator
+from SIDetector import SIDetector
+from IllnessClassifier import IllnessClassifier
+from typing import List, Dict
+
+class UofTearsBot(object):
+    def __init__(self, threshold: float = 0.86, max_history_msgs: int = 50):
+        self.suicidality_detector = SIDetector()
+        self.illness_classifier = IllnessClassifier()
+        self.chatbot = AdviceGenerator()
+        self.history: List[Dict[str, str]] = []
+        self.FLAG = False   # suicidal crisis flag
+        self.threshold = threshold
+        self.max_history_msgs = max_history_msgs
+
+    def safety_check(self, user_text: str):
+        suicidal, confidence = self.suicidality_detector.forward(user_text)
+        self.FLAG = (suicidal and confidence >= self.threshold)
+        if self.FLAG:
+            return "suicidal"
+        disorder, conf = self.illness_classifier.forward(user_text)
+        return disorder
+
+    def userCrisis(self) -> str:
+        return (
+            "I'm really sorry you're feeling this way. Your safety matters.\n"
+            "If you are in immediate danger, please call your local emergency number now.\n"
+            "US & Canada: dial **988** (24/7)\n"
+            "International: call your local emergency services.\n"
+            "If you can, reach out to someone you trust and let them know what’s going on.\n"
+            "We can also talk through coping strategies together."
+        )
+
+    def _prune_history(self) -> List[Dict[str, str]]:
+        """Keep only the last `max_history_msgs` turns."""
+        return self.history[-self.max_history_msgs :]
+
+    def converse(self, user_text: str) -> str:
+        disorder = self.safety_check(user_text)
+
+        # store user text into history
+        self.history.append({"role": "user", "content": user_text})
+
+        if self.FLAG:
+            # crisis flow: respond with fixed crisis message only
+            crisis_msg = self.userCrisis()
+            self.history.append({"role": "assistant", "content": crisis_msg})
+            return crisis_msg
+
+        # normal advice generation
+        pruned_history = self._prune_history()
+        advice = self.chatbot.generate_advice(
+            disorder=disorder,
+            user_text=user_text,
+            history=pruned_history,
+        )['text']
+
+        # add bot response to history
+        self.history.append({"role": "assistant", "content": advice})
+        return advice

app.py

@@ -0,0 +1,68 @@
+import os
+from typing import List, Dict
+import logging
+import dotenv
+
+import torch
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse, HTMLResponse
+from pydantic import BaseModel
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download, login
+
+import uvicorn
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    AutoModelForCausalLM,
+    pipeline
+)
+
+from UofTearsBot import UofTearsBot
+
+MODEL_REPO="bartowski/Mistral-7B-Instruct-v0.3-GGUF"
+MODEL_FILE="Mistral-7B-Instruct-v0.3-Q4_K_M.gguf"
+CHAT_FORMAT="mistral-instruct"
+
+dotenv.load_dotenv()
+login(token=os.getenv("HF_TOKEN"))
+
+MODEL_PATH = hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename=MODEL_FILE,
+    local_dir="/tmp/models",
+    local_dir_use_symlinks=False,
+)
+
+llm = Llama(
+    model_path=MODEL_PATH,
+    n_ctx=int(os.getenv("N_CTX", "4096")),
+    n_threads=os.cpu_count() or 4,
+    n_batch=int(os.getenv("N_BATCH", "256")),
+    chat_format=CHAT_FORMAT,
+)
+
+# Start the FastAPI app
+app = FastAPI()
+chatbots: Dict[str, UofTearsBot] = {}
+
+class ChatRequest(BaseModel):
+    user_id: str
+    user_text: str
+
+@app.post("/chat")
+async def chat(request: ChatRequest):
+    if request.user_id not in chatbots:
+        chatbots[request.user_id] = UofTearsBot(llm)
+    current_bot = chatbots[request.user_id]
+    response = current_bot.converse(request.user_text)
+    return JSONResponse(content={"response": response, "history": current_bot.history})
+
+@app.get("/", response_class=HTMLResponse)
+async def home():
+    return "<h1>App is running 🚀</h1>"
+    
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860) # huggingface port