from transformers import pipeline 
from flask import Flask, jsonify, request, os

app = Flask(__name__)

nlp = pipeline("fill-mask", model=os.environ['MODEL_NAME'])

@app.route('/predict', methods=['POST'])
def predict():
    text = request.json['text']
    result = nlp(text)
    return jsonify(result[0]['sequence'])

if __name__ == '__main__':
    app.run(port=os.environ['API_PORT'])