File size: 729 Bytes
a65ab2b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
from flask import Flask, request, jsonify
from transformers import pipeline
# Initialize the Hugging Face model pipeline
model_name = "distilbert-base-uncased" # Change to the model you need
nlp = pipeline("text-classification", model=model_name)
# Create Flask app
app = Flask(__name__)
@app.route("/predict", methods=["POST"])
def predict():
# Get the data from the request
data = request.get_json()
text = data.get("text", "")
if not text:
return jsonify({"error": "Text input is required"}), 400
# Use Hugging Face model for inference
result = nlp(text)
return jsonify({"prediction": result})
if __name__ == "__main__":
# Run the app
app.run(host="0.0.0.0", port=5000)
|