Spaces:
Runtime error
Runtime error
| # 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) | |
| def predict(): | |
| user_input = request.form['user_input'] | |
| result = predict_depression(user_input) | |
| return result | |
| app.run(debug=True) | |