Spaces:
Runtime error
Runtime error
File size: 1,285 Bytes
9923c53 c428353 9923c53 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
# doing the required imports here
import tensorflow as tf
from transformers import BertTokenizer, TFBertForSequenceClassification
from huggingface_hub import login
from flask import Flask, request
from flask_cors import CORS
# login token for HuggingFace
login(os.getenv('huggingFaceToken'))
model_name = "mental/mental-bert-base-uncased"
# here loading tokenizer
tokenizer = BertTokenizer.from_pretrained(model_name)
# loading bert model --> mental-bert-base-uncase
model = TFBertForSequenceClassification.from_pretrained(model_name, num_labels=3, from_pt=True)
# path to the stored .h5 model
model.load_weights("./mentalbert_model.h5")
# numeric mapping for the label
id2label = {0: "Highly Depressed", 1: "Moderately Depressed", 2: "Not Depressed"}
# predict function that accepts the input text
def predict_depression(text):
inputs = tokenizer(text, return_tensors="tf", padding=True, truncation=True, max_length=64)
logits = model(**inputs).logits
pred_id = tf.argmax(logits, axis=1).numpy()[0]
return id2label[pred_id]
app = Flask(__name__)
CORS(app)
@app.route('/')
@app.route('/predict', methods=['POST'])
def predict():
user_input = request.form['user_input']
result = predict_depression(user_input)
return result
app.run(debug=True)
|