File size: 3,119 Bytes
7711d44 | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | from flask import Flask, request, jsonify
from transformers import pipeline
import logging
# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
app = Flask(__name__)
# Load the sentiment analysis pipeline from Hugging Face
# This will download the model if not already cached
logging.info("Loading sentiment analysis pipeline...")
# For the purpose of this app, we assume the model to be loaded will be 'Arvind111/bert_sentiment-analysis-model' if pushed.
# For local testing, you might use a general sentiment model, but the task specifies using the pushed model.
# Placeholder for the actual model to be loaded from HF. For now, using a general one for structure.
# In a real deployment, we'd replace 'distilbert-base-uncased-finetuned-sst-2-english' with 'Arvind111/bert_sentiment-analysis-model'
# once it's pushed and available.
classifier = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
logging.info("Sentiment analysis pipeline loaded.")
# Define the label mapping as per the trained model (0=Sad, 1=Happy, 2=Neutral)
# The simpletransformers model outputs 0, 1, 2 directly, but the transformers pipeline might output LABEL_0, LABEL_1, etc.
# We need to map these to our desired labels. The simpletransformers model's labels were 0: Sad, 1: happy, 2: Neutral
# If the pipeline returns LABEL_0, LABEL_1, LABEL_2 based on the fine-tuned model, we map accordingly.
# Assuming LABEL_0 -> Sad, LABEL_1 -> Happy, LABEL_2 -> Neutral based on the training data order.
label_mapping = {
'LABEL_0': 'Sad',
'LABEL_1': 'Happy',
'LABEL_2': 'Neutral'
}
@app.route('/')
def home():
return "Welcome to the Emotion Prediction API! Use the /predict endpoint."
@app.route('/predict', methods=['POST'])
def predict():
data = request.get_json(force=True)
text = data.get('text', '')
if not text:
logging.warning("Received empty text for prediction.")
return jsonify({'error': 'No text provided for prediction'}), 400
try:
# The pipeline returns a list of dictionaries, e.g., [{'label': 'LABEL_1', 'score': 0.999}]
prediction_result = classifier(text)[0]
predicted_label_code = prediction_result['label']
score = prediction_result['score']
# Map the Hugging Face label code to our custom emotion label
emotion = label_mapping.get(predicted_label_code, 'Unknown')
response = {
'text': text,
'predicted_emotion': emotion,
'confidence': score
}
logging.info(f"Prediction made for text: '{text[:50]}' - Emotion: {emotion}, Confidence: {score:.4f}")
return jsonify(response)
except Exception as e:
logging.error(f"Error during prediction for text: '{text[:50]}' - Error: {e}", exc_info=True)
return jsonify({'error': str(e)}), 500
if __name__ == '__main__':
# Use a default port or get it from environment variables for deployment
port = 5000
logging.info(f"Starting Flask app on port {port}")
app.run(host='0.0.0.0', port=port)
|