File size: 6,307 Bytes
26c8f44
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
app.py
======
FastAPI application for comment moderation.

Endpoints:
  POST /moderate        β€” Binary moderation (flagged/safe)
  POST /moderate/detail β€” Detailed per-category scores
  GET  /health          β€” Health check

Usage:
    uvicorn api.app:app --host 0.0.0.0 --port 8000
    uvicorn api.app:app --host 0.0.0.0 --port 8000 --reload  # dev mode
"""

import logging
import sys
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware

# Add project root to path for imports
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from api.middleware import TimingMiddleware
from api.schemas import (
    ErrorResponse,
    HealthResponse,
    ModerateDetailResponse,
    ModerateRequest,
    ModerateResponse,
)

# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
)
logger = logging.getLogger(__name__)

# Global moderator instance (loaded at startup)
moderator = None


# ---------------------------------------------------------------------------
# Application lifespan (startup/shutdown)
# ---------------------------------------------------------------------------

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load the moderation model on startup."""
    global moderator

    logger.info("πŸš€ Starting Comment Moderation API...")

    try:
        from inference.predictor import CommentModerator
        moderator = CommentModerator()
        logger.info("βœ… Model loaded successfully!")
    except Exception as e:
        logger.error(f"❌ Failed to load model: {e}")
        logger.error("The API will start but moderation endpoints will return errors.")
        logger.error("Make sure you've trained the model and exported to ONNX:")
        logger.error("  python training/train.py")
        logger.error("  python inference/export_onnx.py")

    yield

    # Cleanup
    logger.info("πŸ‘‹ Shutting down...")
    moderator = None


# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------

app = FastAPI(
    title="Comment Moderation API",
    description=(
        "AI-powered comment moderation system supporting English and Hinglish. "
        "Uses fine-tuned XLM-RoBERTa with ONNX inference for <20ms latency."
    ),
    version="1.0.0",
    lifespan=lifespan,
    docs_url="/docs",
    redoc_url="/redoc",
)

# --- Middleware ---
app.add_middleware(TimingMiddleware)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Restrict in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------

def _check_moderator():
    """Raise 503 if model is not loaded."""
    if moderator is None:
        raise HTTPException(
            status_code=503,
            detail="Model not loaded. Train and export the model first.",
        )


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@app.get(
    "/health",
    response_model=HealthResponse,
    tags=["System"],
    summary="Health check",
)
async def health():
    """Check if the service is running and the model is loaded."""
    return HealthResponse(
        status="healthy" if moderator is not None else "degraded",
        model="xlm-roberta-base-moderation",
        version="1.0.0",
    )


@app.post(
    "/moderate",
    response_model=ModerateResponse,
    tags=["Moderation"],
    summary="Moderate a comment (binary)",
    responses={
        200: {
            "description": "Moderation result",
            "content": {
                "application/json": {
                    "examples": {
                        "flagged": {
                            "summary": "Harmful comment",
                            "value": {"flagged": True},
                        },
                        "safe": {
                            "summary": "Safe comment",
                            "value": {"flagged": False},
                        },
                    }
                }
            },
        },
        503: {"model": ErrorResponse, "description": "Model not loaded"},
    },
)
async def moderate(request: ModerateRequest):
    """
    Classify a comment as harmful or safe.

    Returns `{"flagged": true}` if the comment is harmful,
    `{"flagged": false}` otherwise.
    """
    _check_moderator()
    result = moderator.moderate(request.comment)
    return ModerateResponse(**result)


@app.post(
    "/moderate/detail",
    response_model=ModerateDetailResponse,
    tags=["Moderation"],
    summary="Moderate a comment (detailed)",
    responses={
        200: {"description": "Detailed moderation result"},
        503: {"model": ErrorResponse, "description": "Model not loaded"},
    },
)
async def moderate_detail(request: ModerateRequest):
    """
    Classify a comment with detailed per-category probability scores.

    Returns flagged status, confidence, top harmful category, and
    individual category probabilities.
    """
    _check_moderator()
    result = moderator.predict(request.comment)
    return ModerateDetailResponse(
        flagged=result["flagged"],
        confidence=result["confidence"],
        top_category=result["top_category"],
        categories=result["categories"],
        threshold=moderator.threshold,
        latency_ms=result["latency_ms"],
    )


# ---------------------------------------------------------------------------
# Run directly
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "api.app:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        log_level="info",
    )