Meshyboi/DL-GenAI-Project

1from fastapi import FastAPI, HTTPException2from pydantic import BaseModel3from typing import Dict, List4import os5import tensorflow as tf6import numpy as np7from transformers import RobertaTokenizer, TFRobertaModel8from huggingface_hub import hf_hub_download9import uvicorn10 11# Ensure TensorFlow uses tf_keras (required for transformers compatibility)12os.environ["TF_USE_LEGACY_KERAS"] = "1"13# Import tf_keras to ensure it's available for transformers14try:15    import tf_keras16except ImportError:17    raise ImportError("tf-keras package is required. Install it with: pip install tf-keras")18 19# Initialize FastAPI app20app = FastAPI(21    title="Emotion Classification API",22    description="API for emotion classification using RoBERTa model",23    version="1.0.0"24)25 26# Emotion labels27EMOTIONS = ['anger', 'fear', 'joy', 'sadness', 'surprise']28MAX_LEN = 6129 30# Hugging Face model repository31HF_MODEL_ID = "Meshyboi/Multi-Emotion-Classification"32MODEL_FILENAME = "roberta_emotion_model.keras"33 34# Global variables for model and tokenizer35model = None36tokenizer = None37 38# Request/Response models39class PredictionRequest(BaseModel):40    text: str41 42class EmotionScore(BaseModel):43    emotion: str44    score: float45 46class PredictionResponse(BaseModel):47    text: str48    emotions: Dict[str, float]49    detected_emotions: List[str]50 51def load_model():52    """Load the trained model from Hugging Face"""53    global model54    try:55        if model is None:56            # Download the model file from Hugging Face57            print(f"Downloading model file: {MODEL_FILENAME}")58            model_path = hf_hub_download(59                repo_id=HF_MODEL_ID,60                filename=MODEL_FILENAME,61                cache_dir=None  # Use default cache62            )63            print(f"Model downloaded to: {model_path}")64            65            # Define a dummy weighted_binary_crossentropy function for loading66            # (not used during inference since compile=False)67            def weighted_binary_crossentropy(y_true, y_pred):68                # Dummy implementation - not used during inference69                epsilon = tf.keras.backend.epsilon()70                y_pred = tf.clip_by_value(y_pred, epsilon, 1.0 - epsilon)71                bce = -(y_true * tf.math.log(y_pred) + (1.0 - y_true) * tf.math.log(1.0 - y_pred))72                return tf.reduce_mean(bce)73            74            # Provide custom_objects to handle custom loss function and TFRobertaModel75            # TFRobertaModel is needed because the model architecture uses it76            custom_objects = {77                'weighted_binary_crossentropy': weighted_binary_crossentropy,78                'TFRobertaModel': TFRobertaModel79            }80            81            # Load the model using tf_keras directly (not tf.keras) with custom_objects82            # Use safe_mode=False to allow loading custom objects83            # This is needed because the model was saved with tf_keras and a custom loss84            model = tf_keras.models.load_model(85                model_path, 86                compile=False,87                custom_objects=custom_objects,88                safe_mode=False89            )90            print("Model loaded successfully!")91        return model92    except Exception as e:93        raise RuntimeError(f"Error loading model: {str(e)}")94 95def load_tokenizer():96    """Load the tokenizer from Hugging Face"""97    global tokenizer98    try:99        if tokenizer is None:100            # Download tokenizer files from the tokenizer_files subdirectory101            print("Downloading tokenizer files...")102            tokenizer_files = [103                "tokenizer_files/vocab.json",104                "tokenizer_files/merges.txt",105                "tokenizer_files/tokenizer_config.json",106                "tokenizer_files/special_tokens_map.json"107            ]108            109            # Download all tokenizer files110            for file_path in tokenizer_files:111                hf_hub_download(112                    repo_id=HF_MODEL_ID,113                    filename=file_path,114                    cache_dir=None115                )116            117            # Get the snapshot directory path by downloading the model file (already done)118            # or by downloading any file and getting its parent directory119            # The tokenizer files are in tokenizer_files/ subdirectory of the snapshot120            model_path = hf_hub_download(121                repo_id=HF_MODEL_ID,122                filename=MODEL_FILENAME,123                cache_dir=None124            )125            snapshot_dir = os.path.dirname(model_path)126            tokenizer_dir = os.path.join(snapshot_dir, "tokenizer_files")127            128            print(f"Loading tokenizer from: {tokenizer_dir}")129            130            # Load tokenizer from the local tokenizer_files directory131            tokenizer = RobertaTokenizer.from_pretrained(tokenizer_dir)132            print("Tokenizer loaded successfully!")133        return tokenizer134    except Exception as e:135        raise RuntimeError(f"Error loading tokenizer: {str(e)}")136 137def preprocess_text(text: str, tokenizer, max_len: int):138    """Preprocess text for model input"""139    encoded = tokenizer.encode_plus(140        text,141        add_special_tokens=True,142        max_length=max_len,143        padding='max_length',144        truncation=True,145        return_attention_mask=True,146        return_tensors='tf'147    )148    return encoded['input_ids'], encoded['attention_mask']149 150def predict_emotions(text: str, model, tokenizer):151    """Predict emotions for given text"""152    input_ids, attention_mask = preprocess_text(text, tokenizer, MAX_LEN)153    predictions = model.predict([input_ids, attention_mask], verbose=0)154    return predictions[0]155 156@app.on_event("startup")157async def startup_event():158    print(f"Loading model and tokenizer from Hugging Face: {HF_MODEL_ID}")159    # Load resources160    load_model()161    load_tokenizer()162    print("Model and tokenizer loaded successfully from Hugging Face!")163 164@app.get("/")165async def root():166    """Root endpoint"""167    return {168        "message": "Emotion Classification API",169        "version": "1.0.0",170        "endpoints": {171            "predict": "/predict",172            "health": "/health",173            "docs": "/docs"174        }175    }176 177@app.get("/health")178async def health_check():179    """Health check endpoint"""180    return {181        "status": "healthy",182        "model_loaded": model is not None,183        "tokenizer_loaded": tokenizer is not None184    }185 186@app.post("/predict", response_model=PredictionResponse)187async def predict(request: PredictionRequest):188    """189    Predict emotions for the given text190    191    - **text**: Input text to analyze for emotions192    193    Returns:194    - Dictionary with emotion scores and detected emotions195    """196    if not request.text.strip():197        raise HTTPException(status_code=400, detail="Text cannot be empty")198    199    if model is None or tokenizer is None:200        raise HTTPException(status_code=503, detail="Model or tokenizer not loaded")201    202    try:203        predictions = predict_emotions(request.text, model, tokenizer)204        205        # Create emotion scores dictionary206        emotion_scores = {emotion: float(score) for emotion, score in zip(EMOTIONS, predictions)}207        208        # Detect emotions above threshold209        threshold = 0.5210        detected_emotions = [emotion for emotion, score in emotion_scores.items() if score >= threshold]211        212        return PredictionResponse(213            text=request.text,214            emotions=emotion_scores,215            detected_emotions=detected_emotions216        )217    except Exception as e:218        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")219 220if __name__ == "__main__":221    uvicorn.run(app, host="0.0.0.0", port=7860)222 223