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