# 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)