|
|
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) |
|
|
|
