Spaces:
Sleeping
Sleeping
File size: 9,619 Bytes
7e825f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 | import os
import logging
import time
from datetime import datetime
from typing import Dict, List, Optional
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
from contextlib import asynccontextmanager
# Create logs directory if it doesn't exist
os.makedirs('logs', exist_ok=True)
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler('logs/app.log'),
logging.StreamHandler()
]
)
logger = logging.getLogger(__name__)
# Global variables for model and tokenizer
model = None
tokenizer = None
model_loaded = False
model_info = {
"model_name": "songhieng/roberta-phishing-content-detector-5.0",
"loaded_at": None,
"version": "5.0",
"framework": "transformers"
}
class PredictionRequest(BaseModel):
text: str = Field(..., description="Text content to analyze for phishing", min_length=1, max_length=10000)
class PredictionResponse(BaseModel):
text: str
score: float
description: str
processing_time_ms: float
timestamp: str
class HealthResponse(BaseModel):
status: str
model_loaded: bool
timestamp: str
uptime_seconds: float
class BatchPredictionRequest(BaseModel):
texts: List[str] = Field(..., description="List of texts to analyze", max_items=100)
# Application startup and shutdown events
@asynccontextmanager
async def lifespan(app: FastAPI):
# Startup
logger.info("Starting up the application...")
await load_model()
yield
# Shutdown
logger.info("Shutting down the application...")
app = FastAPI(
title="RoBERTa Phishing Content Detector API",
description="MLOps deployment of RoBERTa model for phishing content detection",
version="5.0.0",
lifespan=lifespan
)
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Startup time for uptime calculation
startup_time = time.time()
async def load_model():
"""Load the model and tokenizer"""
global model, tokenizer, model_loaded, model_info
try:
logger.info("Loading model and tokenizer...")
model_path = "models/roberta-phishing-detector"
if not os.path.exists(model_path):
logger.error(f"Model path {model_path} does not exist!")
raise FileNotFoundError(f"Model not found at {model_path}")
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
model = AutoModelForSequenceClassification.from_pretrained(model_path, local_files_only=True)
# Set model to evaluation mode
model.eval()
model_loaded = True
model_info["loaded_at"] = datetime.now().isoformat()
logger.info("Model and tokenizer loaded successfully!")
except Exception as e:
logger.error(f"Failed to load model: {str(e)}")
model_loaded = False
raise e
def predict_phishing(text: str) -> float:
"""Predict if text is phishing content and return a phishing score
A higher score (closer to 1) indicates more likely to be phishing
A lower score (closer to 0) indicates more likely to be legitimate
"""
if not model_loaded:
raise HTTPException(status_code=503, detail="Model not loaded")
try:
# Tokenize the input text
inputs = tokenizer(
text,
truncation=True,
padding=True,
max_length=4096,
return_tensors="pt"
)
# Make prediction
with torch.no_grad():
outputs = model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
# Get the phishing score (class 1 probability)
phishing_score = float(predictions[0][1])
return phishing_score
except Exception as e:
logger.error(f"Prediction error: {str(e)}")
raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")
@app.middleware("http")
async def log_requests(request: Request, call_next):
"""Log all requests"""
start_time = time.time()
response = await call_next(request)
process_time = time.time() - start_time
logger.info(
f"Request: {request.method} {request.url} - "
f"Status: {response.status_code} - "
f"Time: {process_time:.4f}s"
)
return response
@app.get("/", response_model=Dict)
async def root():
"""Root endpoint"""
return {
"message": "RoBERTa Phishing Content Detector API",
"version": "1.0.0",
"model": model_info["model_name"],
"status": "healthy" if model_loaded else "unhealthy"
}
@app.get("/health", response_model=HealthResponse)
async def health_check():
"""Health check endpoint for monitoring"""
uptime = time.time() - startup_time
return HealthResponse(
status="healthy" if model_loaded else "unhealthy",
model_loaded=model_loaded,
timestamp=datetime.now().isoformat(),
uptime_seconds=uptime
)
@app.get("/model/info")
async def model_info_endpoint():
"""Get model information"""
return {
"model_info": model_info,
"model_loaded": model_loaded,
"torch_version": torch.__version__
}
@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
"""Predict if text content is phishing"""
start_time = time.time()
try:
phishing_score = predict_phishing(request.text)
# Generate description based on score
if phishing_score < 0.2:
classification = "Legitimate (Very Low Risk)"
elif phishing_score < 0.4:
classification = "Likely Legitimate (Low Risk)"
elif phishing_score < 0.6:
classification = "Uncertain (Medium Risk)"
elif phishing_score < 0.8:
classification = "Likely Phishing (High Risk)"
else:
classification = "Phishing (Very High Risk)"
description = f"{classification}: Score {phishing_score:.4f} - Lower scores (closer to 0) indicate legitimate content, higher scores (closer to 1) indicate phishing/malicious content"
processing_time = (time.time() - start_time) * 1000 # Convert to milliseconds
response = PredictionResponse(
text=request.text[:100] + "..." if len(request.text) > 100 else request.text,
score=phishing_score,
description=description,
processing_time_ms=round(processing_time, 2),
timestamp=datetime.now().isoformat()
)
logger.info(f"Prediction made: (phishing score: {phishing_score:.4f}, classification: {classification})")
return response
except Exception as e:
logger.error(f"Prediction endpoint error: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/predict/batch")
async def predict_batch(request: BatchPredictionRequest):
"""Batch prediction endpoint"""
start_time = time.time()
try:
results = []
for text in request.texts:
phishing_score = predict_phishing(text)
# Generate classification based on score
if phishing_score < 0.2:
classification = "Legitimate (Very Low Risk)"
elif phishing_score < 0.4:
classification = "Likely Legitimate (Low Risk)"
elif phishing_score < 0.6:
classification = "Uncertain (Medium Risk)"
elif phishing_score < 0.8:
classification = "Likely Phishing (High Risk)"
else:
classification = "Phishing (Very High Risk)"
results.append({
"text": text[:50] + "..." if len(text) > 50 else text,
"score": phishing_score,
"classification": classification
})
processing_time = (time.time() - start_time) * 1000
return {
"results": results,
"total_processed": len(results),
"processing_time_ms": round(processing_time, 2),
"timestamp": datetime.now().isoformat(),
"note": "Lower scores (closer to 0) indicate legitimate content, higher scores (closer to 1) indicate phishing/malicious content"
}
except Exception as e:
logger.error(f"Batch prediction error: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/metrics")
async def metrics():
"""Basic metrics endpoint for monitoring"""
uptime = time.time() - startup_time
return {
"uptime_seconds": uptime,
"model_loaded": model_loaded,
"model_info": model_info,
"memory_usage": "Not implemented", # Could add psutil for real memory usage
"timestamp": datetime.now().isoformat()
}
if __name__ == "__main__":
uvicorn.run(
"main:app",
host="0.0.0.0",
port=8000,
reload=False,
log_level="info"
) |