app.py

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)