beanbox-apis / app.py
johnpaulbin's picture
Create app.py
939c80c
raw
history blame
737 Bytes
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))