Initial commit with project setup and basic structure established.

app/__init__.py

@@ -0,0 +1,6 @@
+"""
+Application package for the low-latency moderation API.
+"""
+
+__all__ = ["config", "models", "services"]
+

app/config.py

@@ -0,0 +1,37 @@
+"""
+Configuration helpers for the moderation service.
+"""
+
+from dataclasses import dataclass
+from functools import lru_cache
+import os
+
+
+@dataclass(frozen=True)
+class Settings:
+    """Simple immutable settings object."""
+
+    model_name: str = "martin-ha/toxic-comment-model"
+    negative_threshold: float = 0.90
+    api_title: str = "Low-Latency Moderation API"
+    api_description: str = (
+        "Intercepts chat messages and blocks high-confidence toxic content in under 150ms."
+    )
+    api_version: str = "1.0.0"
+
+
+@lru_cache(maxsize=1)
+def get_settings() -> Settings:
+    """
+    Load settings once, allowing environment variables to override defaults.
+
+    Environment variables:
+        MODEL_NAME: Hugging Face model identifier.
+        NEGATIVE_THRESHOLD: Float between 0 and 1 for blocking TOXIC messages.
+    """
+
+    model_name = os.getenv("MODEL_NAME", Settings.model_name)
+    negative_threshold = float(os.getenv("NEGATIVE_THRESHOLD", Settings.negative_threshold))
+
+    return Settings(model_name=model_name, negative_threshold=negative_threshold)
+

app/main.py

@@ -0,0 +1,50 @@
+"""
+FastAPI application entry point for the low-latency moderation API.
+"""
+
+from fastapi import FastAPI, HTTPException
+
+from .config import get_settings
+from .models import MessagePayload, ModerationResponse
+from . import services
+
+settings = get_settings()
+
+app = FastAPI(
+    title=settings.api_title,
+    description=settings.api_description,
+    version=settings.api_version,
+)
+
+
+@app.on_event("startup")
+async def warm_model_cache() -> None:
+    """Load the toxicity classification pipeline during startup to avoid request-time latency."""
+
+    services.get_toxicity_pipeline()
+
+
+@app.post("/api/check-message", response_model=ModerationResponse)
+async def check_message(payload: MessagePayload) -> ModerationResponse:
+    """
+    Classify the incoming text and block highly confident toxic content.
+    """
+
+    result = services.analyze_text(payload.text)
+
+    if result.label == "TOXIC" and result.confidence >= settings.negative_threshold:
+        return ModerationResponse(
+                status="rejected",
+                message="Message classified as toxic with high confidence.",
+                label=result.label,
+                confidence=result.confidence,
+            )
+        
+
+    return ModerationResponse(
+        status= "rejected" if result.label == "TOXIC" else "accepted",
+        message="Message passed moderation.",
+        label=result.label,
+        confidence=result.confidence,
+    )
+

app/models.py

@@ -0,0 +1,30 @@
+"""
+Pydantic models shared across the API.
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class MessagePayload(BaseModel):
+    """Incoming chat message payload."""
+
+    text: str = Field(..., min_length=1, description="User-provided message to moderate.")
+
+
+class ToxicityResult(BaseModel):
+    """Normalized output from the toxicity classification model."""
+
+    label: Literal["TOXIC", "NON-TOXIC"]
+    confidence: float = Field(..., ge=0.0, le=1.0)
+
+
+class ModerationResponse(BaseModel):
+    """API response returned to the chat application."""
+
+    status: Literal["accepted", "rejected"]
+    message: str
+    label: Literal["TOXIC", "NON-TOXIC"]
+    confidence: float
+

app/services.py

@@ -0,0 +1,51 @@
+"""
+Service layer for interacting with the Hugging Face toxic comment classification pipeline.
+"""
+
+from functools import lru_cache
+from typing import Callable, List, TypedDict
+
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, TextClassificationPipeline
+
+from .config import get_settings
+from .models import ToxicityResult
+
+
+class PipelineOutput(TypedDict):
+    label: str
+    score: float
+
+
+def _build_pipeline() -> Callable[[str], List[PipelineOutput]]:
+    """
+    Instantiate the Hugging Face TextClassificationPipeline for toxic comment detection.
+    Returns a callable that accepts text and yields label/score dictionaries.
+    """
+
+    settings = get_settings()
+    tokenizer = AutoTokenizer.from_pretrained(settings.model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(settings.model_name)
+    return TextClassificationPipeline(model=model, tokenizer=tokenizer)
+
+
+@lru_cache(maxsize=1)
+def get_toxicity_pipeline() -> Callable[[str], List[PipelineOutput]]:
+    """Return a cached instance of the toxicity classification pipeline."""
+
+    return _build_pipeline()
+
+
+def analyze_text(text: str) -> ToxicityResult:
+    """
+    Run toxicity classification on text and normalize the response.
+    The model outputs labels that we normalize to TOXIC/NON-TOXIC.
+    """
+
+    predictor = get_toxicity_pipeline()
+    result = predictor(text)[0]
+    raw_label = result["label"].upper()
+    confidence = float(result["score"])
+    # print(f"{raw_label}....raw_label, {result}")
+
+    return ToxicityResult(label=raw_label, confidence=confidence)
+

check_model_size.py

@@ -0,0 +1,71 @@
+"""
+Script to check the memory usage of the loaded toxicity classification model.
+"""
+
+import os
+import psutil
+import torch
+
+from app.services import get_toxicity_pipeline
+
+
+def format_bytes(bytes_value: int) -> str:
+    """Convert bytes to human-readable format."""
+    for unit in ['B', 'KB', 'MB', 'GB']:
+        if bytes_value < 1024.0:
+            return f"{bytes_value:.2f} {unit}"
+        bytes_value /= 1024.0
+    return f"{bytes_value:.2f} TB"
+
+
+def main():
+    """Load the model and display memory usage statistics."""
+    
+    process = psutil.Process(os.getpid())
+    
+    # Get baseline memory before loading
+    baseline_memory = process.memory_info().rss
+    print(f"Baseline RAM usage: {format_bytes(baseline_memory)}")
+    print()
+    
+    print("Loading model...")
+    get_toxicity_pipeline()  # Forces cache to load it
+    print("Model loaded!")
+    print()
+    
+    # Get memory after loading
+    final_memory = process.memory_info().rss
+    model_memory = final_memory - baseline_memory
+    
+    print("=" * 50)
+    print("Memory Statistics:")
+    print("=" * 50)
+    print(f"Baseline RAM:     {format_bytes(baseline_memory)}")
+    print(f"Final RAM:        {format_bytes(final_memory)}")
+    print(f"Model RAM:        {format_bytes(model_memory)}")
+    print()
+    
+    # Additional system info
+    print("System Information:")
+    print("=" * 50)
+    total_memory = psutil.virtual_memory().total
+    available_memory = psutil.virtual_memory().available
+    print(f"Total RAM:         {format_bytes(total_memory)}")
+    print(f"Available RAM:     {format_bytes(available_memory)}")
+    print(f"RAM Used:          {format_bytes(total_memory - available_memory)}")
+    print()
+    
+    # PyTorch GPU info if available
+    if torch.cuda.is_available():
+        print("GPU Information:")
+        print("=" * 50)
+        for i in range(torch.cuda.device_count()):
+            gpu_memory = torch.cuda.get_device_properties(i).total_memory
+            print(f"GPU {i}:            {format_bytes(gpu_memory)}")
+    else:
+        print("GPU: Not available (using CPU)")
+
+
+if __name__ == "__main__":
+    main()
+

main.py

@@ -0,0 +1 @@
+from app.main import app

requirements.txt

@@ -0,0 +1,8 @@
+fastapi==0.111.0
+uvicorn[standard]==0.30.1
+transformers==4.44.0
+torch==2.3.1
+pydantic==2.7.1
+pytest==8.3.2
+psutil==5.9.8
+

tests/test_api.py

@@ -0,0 +1,55 @@
+"""
+Lightweight tests for the moderation endpoint.
+"""
+
+from fastapi.testclient import TestClient
+import pytest
+
+from app.main import app
+from app import services
+from app.models import SentimentResult
+
+client = TestClient(app)
+
+
+@pytest.fixture(autouse=True)
+def clear_cache():
+    """Ensure per-test isolation for cached pipeline calls."""
+
+    services.get_sentiment_pipeline.cache_clear()  # type: ignore[attr-defined]
+    yield
+    services.get_sentiment_pipeline.cache_clear()  # type: ignore[attr-defined]
+
+
+def test_rejects_high_confidence_negative(monkeypatch):
+    """Requests should be blocked when the model is confident a message is negative."""
+
+    monkeypatch.setattr(
+        services,
+        "analyze_text",
+        lambda _text: SentimentResult(label="NEGATIVE", confidence=0.93),
+    )
+
+    response = client.post("/api/check-message", json={"text": "You are awful."})
+    assert response.status_code == 400
+    detail = response.json()
+    assert detail["status"] == "rejected"
+    assert detail["label"] == "NEGATIVE"
+
+
+def test_accepts_positive_or_low_confidence(monkeypatch):
+    """Requests should succeed when the message is allowed."""
+
+    monkeypatch.setattr(
+        services,
+        "analyze_text",
+        lambda _text: SentimentResult(label="POSITIVE", confidence=0.52),
+    )
+
+    response = client.post("/api/check-message", json={"text": "Great job!"})
+    payload = response.json()
+
+    assert response.status_code == 200
+    assert payload["status"] == "accepted"
+    assert payload["label"] == "POSITIVE"
+