Update app.py

app.py

@@ -11,11 +11,14 @@ tokenizer.pad_token_id = tokenizer.eos_token_id
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL, low_cpu_mem_usage=True, return_dict=True,torch_dtype=torch.float16,
     device_map="cpu")
 
-def prediction(news):
+def prediction(text):
     # create pipeline
-    clasifer = pipeline("text-generation", tokenizer=tokenizer, model=model, torch_dtype=torch.float16,
+    pipe = pipeline("text-generation", tokenizer=tokenizer, model=model, torch_dtype=torch.float16,
     device_map="cpu",)
 
+    prompt = f"""Classify the text into Normal, Depression, Anxiety, Bipolar, and return the answer as the corresponding mental health disorder label.
+    text: {text}
+    label: """.strip()
     outputs = pipe(prompt, max_new_tokens=2, do_sample=True, temperature=0.1)
     preds = outputs[0]["generated_text"].split("label: ")[-1].strip()