from fastapi import FastAPI, HTTPException from pydantic import BaseModel from typing import Dict, List import os import tensorflow as tf import numpy as np from transformers import RobertaTokenizer, TFRobertaModel from huggingface_hub import hf_hub_download import uvicorn # Ensure TensorFlow uses tf_keras (required for transformers compatibility) os.environ["TF_USE_LEGACY_KERAS"] = "1" # Import tf_keras to ensure it's available for transformers try: import tf_keras except ImportError: raise ImportError("tf-keras package is required. Install it with: pip install tf-keras") # Initialize FastAPI app app = FastAPI( title="Emotion Classification API", description="API for emotion classification using RoBERTa model", version="1.0.0" ) # Emotion labels EMOTIONS = ['anger', 'fear', 'joy', 'sadness', 'surprise'] MAX_LEN = 61 # Hugging Face model repository HF_MODEL_ID = "Meshyboi/Multi-Emotion-Classification" MODEL_FILENAME = "roberta_emotion_model.keras" # Global variables for model and tokenizer model = None tokenizer = None # Request/Response models class PredictionRequest(BaseModel): text: str class EmotionScore(BaseModel): emotion: str score: float class PredictionResponse(BaseModel): text: str emotions: Dict[str, float] detected_emotions: List[str] def load_model(): """Load the trained model from Hugging Face""" global model try: if model is None: # Download the model file from Hugging Face print(f"Downloading model file: {MODEL_FILENAME}") model_path = hf_hub_download( repo_id=HF_MODEL_ID, filename=MODEL_FILENAME, cache_dir=None # Use default cache ) print(f"Model downloaded to: {model_path}") # Define a dummy weighted_binary_crossentropy function for loading # (not used during inference since compile=False) def weighted_binary_crossentropy(y_true, y_pred): # Dummy implementation - not used during inference epsilon = tf.keras.backend.epsilon() y_pred = tf.clip_by_value(y_pred, epsilon, 1.0 - epsilon) bce = -(y_true * tf.math.log(y_pred) + (1.0 - y_true) * tf.math.log(1.0 - y_pred)) return tf.reduce_mean(bce) # Provide custom_objects to handle custom loss function and TFRobertaModel # TFRobertaModel is needed because the model architecture uses it custom_objects = { 'weighted_binary_crossentropy': weighted_binary_crossentropy, 'TFRobertaModel': TFRobertaModel } # Load the model using tf_keras directly (not tf.keras) with custom_objects # Use safe_mode=False to allow loading custom objects # This is needed because the model was saved with tf_keras and a custom loss model = tf_keras.models.load_model( model_path, compile=False, custom_objects=custom_objects, safe_mode=False ) print("Model loaded successfully!") return model except Exception as e: raise RuntimeError(f"Error loading model: {str(e)}") def load_tokenizer(): """Load the tokenizer from Hugging Face""" global tokenizer try: if tokenizer is None: # Download tokenizer files from the tokenizer_files subdirectory print("Downloading tokenizer files...") tokenizer_files = [ "tokenizer_files/vocab.json", "tokenizer_files/merges.txt", "tokenizer_files/tokenizer_config.json", "tokenizer_files/special_tokens_map.json" ] # Download all tokenizer files for file_path in tokenizer_files: hf_hub_download( repo_id=HF_MODEL_ID, filename=file_path, cache_dir=None ) # Get the snapshot directory path by downloading the model file (already done) # or by downloading any file and getting its parent directory # The tokenizer files are in tokenizer_files/ subdirectory of the snapshot model_path = hf_hub_download( repo_id=HF_MODEL_ID, filename=MODEL_FILENAME, cache_dir=None ) snapshot_dir = os.path.dirname(model_path) tokenizer_dir = os.path.join(snapshot_dir, "tokenizer_files") print(f"Loading tokenizer from: {tokenizer_dir}") # Load tokenizer from the local tokenizer_files directory tokenizer = RobertaTokenizer.from_pretrained(tokenizer_dir) print("Tokenizer loaded successfully!") return tokenizer except Exception as e: raise RuntimeError(f"Error loading tokenizer: {str(e)}") def preprocess_text(text: str, tokenizer, max_len: int): """Preprocess text for model input""" encoded = tokenizer.encode_plus( text, add_special_tokens=True, max_length=max_len, padding='max_length', truncation=True, return_attention_mask=True, return_tensors='tf' ) return encoded['input_ids'], encoded['attention_mask'] def predict_emotions(text: str, model, tokenizer): """Predict emotions for given text""" input_ids, attention_mask = preprocess_text(text, tokenizer, MAX_LEN) predictions = model.predict([input_ids, attention_mask], verbose=0) return predictions[0] @app.on_event("startup") async def startup_event(): print(f"Loading model and tokenizer from Hugging Face: {HF_MODEL_ID}") # Load resources load_model() load_tokenizer() print("Model and tokenizer loaded successfully from Hugging Face!") @app.get("/") async def root(): """Root endpoint""" return { "message": "Emotion Classification API", "version": "1.0.0", "endpoints": { "predict": "/predict", "health": "/health", "docs": "/docs" } } @app.get("/health") async def health_check(): """Health check endpoint""" return { "status": "healthy", "model_loaded": model is not None, "tokenizer_loaded": tokenizer is not None } @app.post("/predict", response_model=PredictionResponse) async def predict(request: PredictionRequest): """ Predict emotions for the given text - **text**: Input text to analyze for emotions Returns: - Dictionary with emotion scores and detected emotions """ if not request.text.strip(): raise HTTPException(status_code=400, detail="Text cannot be empty") if model is None or tokenizer is None: raise HTTPException(status_code=503, detail="Model or tokenizer not loaded") try: predictions = predict_emotions(request.text, model, tokenizer) # Create emotion scores dictionary emotion_scores = {emotion: float(score) for emotion, score in zip(EMOTIONS, predictions)} # Detect emotions above threshold threshold = 0.5 detected_emotions = [emotion for emotion, score in emotion_scores.items() if score >= threshold] return PredictionResponse( text=request.text, emotions=emotion_scores, detected_emotions=detected_emotions ) except Exception as e: raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}") if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=7860)