Update main.py

main.py

@@ -1,27 +1,22 @@
-from ctransformers import AutoModelForCausalLM
-from fastapi import FastAPI, Form
+from fastapi import FastAPI
 from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 
-#Model loading
-llm = AutoModelForCausalLM.from_pretrained("TherapyBeagle-11B-v1-Q2_K.gguf",
-model_type='llama',
-max_new_tokens = 1096,
-threads = 3,
-)
-   
+# Load the depression detection model
+tokenizer = AutoTokenizer.from_pretrained("paulagarciaserrano/roberta-depression-detection")
+model = AutoModelForSequenceClassification.from_pretrained("paulagarciaserrano/roberta-depression-detection")
+depression_pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
 
-#Pydantic object
-class validation(BaseModel):
-    prompt: str
-    
-#Fast API
+# Pydantic object for validation
+class DepressionInput(BaseModel):
+    text: str
+
+# FastAPI app
 app = FastAPI()
 
-#Zephyr completion
-@app.post("/llm_on_cpu")
-async def stream(item: validation):
-    system_prompt = 'Below is an instruction that describes a task. Write a response that appropriately completes the request.'
-    E_INST = "</s>"
-    user, assistant = "<|user|>", "<|assistant|>"
-    prompt = f"{system_prompt}{E_INST}\n{user}\n{item.prompt.strip()}{E_INST}\n{assistant}\n"
-    return llm(prompt)
+# Endpoint for depression detection
+@app.post("/depression_detection")
+async def detect_depression(item: DepressionInput):
+    result = depression_pipe(item.text)
+    return result
+