from flask import Flask, request, jsonify
import asyncio
from hypercorn.asyncio import serve
from hypercorn.config import Config
from setfit import SetFitModel

app = Flask(__name__)

model = SetFitModel.from_pretrained("johnpaulbin/beanbox-toxic")


@app.route('/infer', methods=['POST'])
def translate():
    data = request.get_json()
    result = model.predict_proba([data['text']])

    if result[0][0] > result[0][1]:
        result = "FALSE"
    else:
        result = "TRUE"

    return jsonify(result)

# Define more routes for other operations like download_model, etc.
if __name__ == "__main__":
   config = Config()
   config.bind = ["0.0.0.0:7860"]  # You can specify the host and port here
   asyncio.run(serve(app, config))