app.py

import os
os.environ["HF_HOME"] = "/app/cache"
os.environ["TORCH_HOME"] = "/app/cache"

import torch
import sys 
import torch.nn as nn
from torchvision import transforms
from torchvision.models import inception_v3, Inception_V3_Weights
from transformers import BertTokenizerFast, BertModel
from PIL import Image
from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse
import traceback
from io import BytesIO

# ----------------------------
# Configurations
# ----------------------------
app = FastAPI(title="Multimodal Hate Speech Detection")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

CACHE_DIR = "/app/cache"
MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "BERT_Inception_multimodal.pt")

# Image preprocessing
image_transform = transforms.Compose([
    transforms.Resize((299, 299)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# ----------------------------
# Model Definition
# ----------------------------
class TensorFusionMultimodalModel(nn.Module):
    def __init__(self, bert_model):
        super(TensorFusionMultimodalModel, self).__init__()
        self.bert = bert_model
        self.text_fc = nn.Linear(self.bert.config.hidden_size, 256)

        # Image branch (InceptionV3 backbone)
        weights = Inception_V3_Weights.IMAGENET1K_V1
        inception = inception_v3(weights=None, aux_logits=True)  # Do NOT auto-download
        state_dict = weights.get_state_dict(progress=False)  # Load from local cache
        inception.load_state_dict(state_dict)

        for param in inception.parameters():
            param.requires_grad = False

        self.cnn_backbone = nn.Sequential(*list(inception.children())[:-1])  # Remove FC layer
        self.image_fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(2048, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.5)
        )

        # Fusion layer
        self.fusion_dim = 256 * 256 + 256 + 256 + 1
        self.fusion_fc = nn.Sequential(
            nn.Linear(self.fusion_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

    def tensor_fusion(self, text_feat, img_feat):
        batch_size = text_feat.size(0)
        outer = torch.bmm(text_feat.unsqueeze(2), img_feat.unsqueeze(1)).view(batch_size, -1)
        fusion = torch.cat([outer, text_feat, img_feat, torch.ones(batch_size, 1).to(text_feat.device)], dim=1)
        return fusion

    def forward(self, input_ids, attention_mask, images):
        text_output = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        text_feat = self.text_fc(text_output.pooler_output)

        img_feat = self.cnn_backbone(images).flatten(1)
        img_feat = self.image_fc(img_feat)

        fused_feat = self.tensor_fusion(text_feat, img_feat)
        return self.fusion_fc(fused_feat)

# ----------------------------
# Load Model on Startup
# ----------------------------
@app.on_event("startup")
def load_model():
    global model, tokenizer

    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"[ERROR] Model file not found: {MODEL_PATH}")

    tokenizer = BertTokenizerFast.from_pretrained("bert-base-multilingual-cased", cache_dir=CACHE_DIR)

    # Allow pickle resolution
    sys.modules['__main__'] = sys.modules['app']

    # Load the full model
    loaded = torch.load(MODEL_PATH, map_location=device)

    # ✅ If loaded object still has `forward` via DataParallel, unwrap forcibly
    if hasattr(loaded, "module"):  # Works even if type doesn't match
        loaded = loaded.module

    # ✅ Move all parameters and buffers to device
    loaded = loaded.to(device)

    # Set to eval
    loaded.eval()

    for p in loaded.parameters():
        p.requires_grad = False

    model = loaded
    print(f"[INFO] Model loaded successfully on {device}")

# @app.on_event("startup")
# def load_model():
#     global model, tokenizer

#     if not os.path.exists(MODEL_PATH):
#         raise FileNotFoundError(f"[ERROR] Model file not found: {MODEL_PATH}")

#     # Load tokenizer & BERT
#     tokenizer = BertTokenizerFast.from_pretrained("bert-base-multilingual-cased", cache_dir=CACHE_DIR)
#     bert_model = BertModel.from_pretrained("bert-base-multilingual-cased", cache_dir=CACHE_DIR)

#     # Trick for pickle to resolve class path
#     import sys
#     sys.modules['__main__'] = sys.modules['app']

#     # Load the full model
#     model = torch.load(MODEL_PATH, map_location=device)
#     model.to(device)
#     model.eval()

#     print("[INFO] Model and tokenizer loaded successfully!")





# ----------------------------
# API Endpoints
# ----------------------------
@app.get("/")
async def root():
    return {"message": "API is running. Use POST /predict"}


@app.post("/predict")
async def predict(text: str = Form(...), file: UploadFile = File(...)):
    try:
        if 'model' not in globals() or model is None:
            return JSONResponse(content={"error": "Model not loaded"}, status_code=500)
        if 'tokenizer' not in globals() or tokenizer is None:
            return JSONResponse(content={"error": "Tokenizer not loaded"}, status_code=500)

        # Tokenize text
        encoding = tokenizer(
            text,
            add_special_tokens=True,
            max_length=125,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        input_ids = encoding['input_ids'].to(device, non_blocking=True)
        attention_mask = encoding['attention_mask'].to(device, non_blocking=True)

        # Process image
        contents = await file.read()
        if not contents:
            return JSONResponse(content={"error": "Empty file received"}, status_code=400)

        try:
            img = Image.open(BytesIO(contents)).convert("RGB")
        except Exception:
            return JSONResponse(content={"error": "Invalid image file"}, status_code=400)

        img_tensor = image_transform(img).unsqueeze(0).to(device, non_blocking=True)

        # Predict
        with torch.inference_mode():
            # (Optional) autocast for GPU float16 speedup; safe to leave on CPU too
            # On CPU autocast('cpu') is available in newer torch; we keep it simple:
            output = model(input_ids, attention_mask, img_tensor)

            # output should be shape [1, 1]; get scalar
            prob = float(output.squeeze().item())
            label = "Hate Speech" if prob >= 0.5 else "Non-Hate Speech"

        return {"prediction": label, "confidence": round(prob, 4)}

    except Exception:
        print("Prediction Error:", traceback.format_exc())
        return JSONResponse(content={"error": "Internal server error"}, status_code=500)

# @app.post("/predict")
# async def predict(text: str = Form(...), file: UploadFile = File(...)):
#     try:
#         # Tokenize text
#         encoding = tokenizer(
#             text,
#             add_special_tokens=True,
#             max_length=125,
#             padding='max_length',
#             truncation=True,
#             return_tensors='pt'
#         )
#         input_ids = encoding['input_ids'].to(device)
#         attention_mask = encoding['attention_mask'].to(device)

#         # Process image
#         contents = await file.read()
#         img = Image.open(BytesIO(contents)).convert("RGB")
#         img_tensor = image_transform(img).unsqueeze(0).to(device)

#         # Predict
#         with torch.no_grad():
#             output = model(input_ids, attention_mask, img_tensor)
#             prob = output.item()
#             label = "Hate" if prob >= 0.5 else "No Hate"

#         return {"prediction": label, "confidence": round(prob, 4)}


    
#     except Exception as e:
#         print("Prediction Error:", traceback.format_exc())
#         return JSONResponse(content={"error": str(e)}, status_code=500)


# ----------------------------
# Run server (for local dev)
# ----------------------------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=False)