from flask import Flask, request, jsonify
import torch

from minilm_keyword_pred import predict_on_text

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
print('dev', device)

app = Flask(__name__)

@app.route('/get_keywords', methods=['POST'])
def get_keywords():
  text  = request.json.get('text', '')
  if not text: return jsonify({"error": "The 'text' field is required."}), 400 
  groups = predict_on_text(text)['predicted_groups_with_scores']
  groups = [{'keyword': x[0], 'score': float(x[1])} for x in groups]
  return jsonify(groups)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=14005, debug=True, use_reloader=False)