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)