app.py

import os
import time
import threading
import torch
import json
from typing import List, Dict, Any, Optional
from fastapi import FastAPI, HTTPException, Depends, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel, Field
from dotenv import load_dotenv
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import uvicorn

# Load environment variables
load_dotenv()

# Create necessary directories
os.makedirs("templates", exist_ok=True)
os.makedirs("static", exist_ok=True)

# Model configuration
MODEL_REPO = "daniel-dona/gemma-3-270m-it"
LOCAL_DIR = os.path.join(os.getcwd(), "local_model")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
os.environ.setdefault("OMP_NUM_THREADS", str(os.cpu_count() or 2))
os.environ.setdefault("MKL_NUM_THREADS", os.environ["OMP_NUM_THREADS"])
os.environ.setdefault("OMP_PROC_BIND", "TRUE")
torch.set_num_threads(int(os.environ["OMP_NUM_THREADS"]))
torch.set_num_interop_threads(1)
torch.set_float32_matmul_precision("high")

# Initialize FastAPI app
app = FastAPI(title="AI Content Moderator API", description="Advanced content moderation API powered by AI")

# Mount static files and templates
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

# Ensure model is downloaded
def ensure_local_model(repo_id: str, local_dir: str, tries: int = 3, sleep_s: float = 3.0) -> str:
    os.makedirs(local_dir, exist_ok=True)
    for i in range(tries):
        try:
            snapshot_download(
                repo_id=repo_id,
                local_dir=local_dir,
                local_dir_use_symlinks=False,
                resume_download=True,
                allow_patterns=["*.json", "*.model", "*.safetensors", "*.bin", "*.txt", "*.py"]
            )
            return local_dir
        except Exception:
            if i == tries - 1:
                raise
            time.sleep(sleep_s * (2 ** i))
    return local_dir

# Load model and tokenizer
print("Loading model...")
model_path = ensure_local_model(MODEL_REPO, LOCAL_DIR)
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)

# Define chat template
gemma_chat_template_simplified = (
    "{% for message in messages %}"
        "{% if message['role'] == 'user' %}"
            "{{ '<start_of_turn>user\\n' + message['content'] | trim + '<end_of_turn>\\n' }}"
        "{% elif message['role'] == 'assistant' %}"
            "{{ '<start_of_turn>model\\n' + message['content'] | trim + '<end_of_turn>\\n' }}"
        "{% endif %}"
    "{% endfor %}"
    "{% if add_generation_prompt %}"
        "{{ '<start_of_turn>model\\n' }}"
    "{% endif %}"
)

if tokenizer.chat_template is None:
    tokenizer.chat_template = gemma_chat_template_simplified

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    local_files_only=True,
    torch_dtype=torch.float32,
    device_map=None
)
model.eval()

# System prompt for moderation
MODERATION_SYSTEM_PROMPT = (
    "You are a multilingual content moderation classifier. "
    "You MUST respond with exactly one lowercase letter: 's' for safe, 'u' for unsafe. "
    "No explanations, no punctuation, no extra words. "
    "If the message contains hate speech, harassment, sexual content involving minors, "
    "extreme violence, self-harm encouragement, or other unsafe material, respond 'u'. "
    "Otherwise respond 's'."
)

# API Models
class ModerationRequest(BaseModel):
    input: str = Field(..., description="Text to moderate")
    model: Optional[str] = Field("gemma-3-270m-it", description="Model to use for moderation")

class BatchModerationRequest(BaseModel):
    inputs: List[str] = Field(..., description="List of texts to moderate")
    model: Optional[str] = Field("gemma-3-270m-it", description="Model to use for moderation")

class ModerationResponse(BaseModel):
    id: str
    object: str
    created: int
    model: str
    results: List[Dict[str, Any]]

# Helper functions
def build_prompt(message, max_ctx_tokens=128):
    full_user_message = f"{MODERATION_SYSTEM_PROMPT}\n\nUser input: '{message}'"
    messages = [{"role": "user", "content": full_user_message}]
    
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    while len(tokenizer(text, add_special_tokens=False).input_ids) > max_ctx_tokens and len(full_user_message) > 100:
        full_user_message = full_user_message[:-50]
        messages[0]['content'] = full_user_message
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
    return text

def enforce_s_u(text: str) -> str:
    text_lower = text.strip().lower()
    if "u" in text_lower and "s" not in text_lower:
        return "u"
    if "unsafe" in text_lower:
        return "u"
    return "s"

def classify_text(message, max_tokens=3, temperature=0.1, top_p=0.95):
    if not message.strip():
        return {
            "classification": "s",
            "label": "SAFE",
            "description": "Content appears to be safe and appropriate.",
            "tokens_per_second": 0,
            "processing_time": 0
        }
    
    text = build_prompt(message)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    do_sample = bool(temperature and temperature > 0.0)
    gen_kwargs = dict(
        max_new_tokens=max_tokens,
        do_sample=do_sample,
        top_p=top_p,
        temperature=temperature if do_sample else None,
        use_cache=True,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id
    )
    
    try:
        streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True, skip_prompt=True)
    except TypeError:
        streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
    
    thread = threading.Thread(
        target=model.generate,
        kwargs={**inputs, **{k: v for k, v in gen_kwargs.items() if v is not None}, "streamer": streamer}
    )
    
    partial_text = ""
    token_count = 0
    start_time = None
    
    with torch.inference_mode():
        thread.start()
        try:
            for chunk in streamer:
                if start_time is None:
                    start_time = time.time()
                partial_text += chunk
                token_count += 1
        finally:
            thread.join()
    
    final_label = enforce_s_u(partial_text)
    end_time = time.time() if start_time else time.time()
    duration = max(1e-6, end_time - start_time)
    tps = token_count / duration if duration > 0 else 0.0
    
    if final_label == "s":
        label = "SAFE"
        description = "Content appears to be safe and appropriate."
    else:
        label = "UNSAFE"
        description = "Content may contain inappropriate or harmful material."
    
    return {
        "classification": final_label,
        "label": label,
        "description": description,
        "tokens_per_second": tps,
        "processing_time": duration
    }

# API Key validation
def get_api_key(request: Request):
    api_key = request.headers.get("Authorization") or request.query_params.get("api_key")
    if not api_key:
        raise HTTPException(status_code=401, detail="API key required")
    
    # Remove "Bearer " prefix if present
    if api_key.startswith("Bearer "):
        api_key = api_key[7:]
    
    # Validate against environment variable
    env_api_key = os.getenv("API_KEY")
    if not env_api_key or api_key != env_api_key:
        raise HTTPException(status_code=401, detail="Invalid API key")
    
    return api_key

# API Endpoints
@app.get("/", response_class=HTMLResponse)
async def get_home(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/v1/moderations", response_model=ModerationResponse)
async def moderate_content(
    request: ModerationRequest,
    api_key: str = Depends(get_api_key)
):
    result = classify_text(request.input)
    
    response_data = {
        "id": f"modr_{int(time.time())}",
        "object": "moderation",
        "created": int(time.time()),
        "model": request.model,
        "results": [
            {
                "flagged": result["classification"] == "u",
                "categories": {
                    "hate": result["classification"] == "u",
                    "hate/threatening": result["classification"] == "u",
                    "harassment": result["classification"] == "u",
                    "harassment/threatening": result["classification"] == "u",
                    "self-harm": result["classification"] == "u",
                    "self-harm/intent": result["classification"] == "u",
                    "self-harm/instructions": result["classification"] == "u",
                    "sexual": result["classification"] == "u",
                    "sexual/minors": result["classification"] == "u",
                    "violence": result["classification"] == "u",
                    "violence/graphic": result["classification"] == "u"
                },
                "category_scores": {
                    "hate": 0.9 if result["classification"] == "u" else 0.1,
                    "hate/threatening": 0.9 if result["classification"] == "u" else 0.1,
                    "harassment": 0.9 if result["classification"] == "u" else 0.1,
                    "harassment/threatening": 0.9 if result["classification"] == "u" else 0.1,
                    "self-harm": 0.9 if result["classification"] == "u" else 0.1,
                    "self-harm/intent": 0.9 if result["classification"] == "u" else 0.1,
                    "self-harm/instructions": 0.9 if result["classification"] == "u" else 0.1,
                    "sexual": 0.9 if result["classification"] == "u" else 0.1,
                    "sexual/minors": 0.9 if result["classification"] == "u" else 0.1,
                    "violence": 0.9 if result["classification"] == "u" else 0.1,
                    "violence/graphic": 0.9 if result["classification"] == "u" else 0.1
                },
                "text": request.input
            }
        ]
    }
    
    return response_data

@app.post("/v1/moderations/batch", response_model=ModerationResponse)
async def moderate_content_batch(
    request: BatchModerationRequest,
    api_key: str = Depends(get_api_key)
):
    results = []
    
    for text in request.inputs:
        result = classify_text(text)
        results.append({
            "flagged": result["classification"] == "u",
            "categories": {
                "hate": result["classification"] == "u",
                "hate/threatening": result["classification"] == "u",
                "harassment": result["classification"] == "u",
                "harassment/threatening": result["classification"] == "u",
                "self-harm": result["classification"] == "u",
                "self-harm/intent": result["classification"] == "u",
                "self-harm/instructions": result["classification"] == "u",
                "sexual": result["classification"] == "u",
                "sexual/minors": result["classification"] == "u",
                "violence": result["classification"] == "u",
                "violence/graphic": result["classification"] == "u"
            },
            "category_scores": {
                "hate": 0.9 if result["classification"] == "u" else 0.1,
                "hate/threatening": 0.9 if result["classification"] == "u" else 0.1,
                "harassment": 0.9 if result["classification"] == "u" else 0.1,
                "harassment/threatening": 0.9 if result["classification"] == "u" else 0.1,
                "self-harm": 0.9 if result["classification"] == "u" else 0.1,
                "self-harm/intent": 0.9 if result["classification"] == "u" else 0.1,
                "self-harm/instructions": 0.9 if result["classification"] == "u" else 0.1,
                "sexual": 0.9 if result["classification"] == "u" else 0.1,
                "sexual/minors": 0.9 if result["classification"] == "u" else 0.1,
                "violence": 0.9 if result["classification"] == "u" else 0.1,
                "violence/graphic": 0.9 if result["classification"] == "u" else 0.1
            },
            "text": text
        })
    
    response_data = {
        "id": f"modr_batch_{int(time.time())}",
        "object": "moderation",
        "created": int(time.time()),
        "model": request.model,
        "results": results
    }
    
    return response_data

# Create the HTML template with Tailwind CSS
with open("templates/index.html", "w") as f:
    f.write("""
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>AI Content Moderator</title>
    <script src="https://cdn.tailwindcss.com"></script>
    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css">
    <style>
        @keyframes float {
            0% { transform: translateY(0px); }
            50% { transform: translateY(-10px); }
            100% { transform: translateY(0px); }
        }
        @keyframes pulse-border {
            0% { box-shadow: 0 0 0 0 rgba(99, 102, 241, 0.7); }
            70% { box-shadow: 0 0 0 10px rgba(99, 102, 241, 0); }
            100% { box-shadow: 0 0 0 0 rgba(99, 102, 241, 0); }
        }
        .float-animation {
            animation: float 3s ease-in-out infinite;
        }
        .pulse-border {
            animation: pulse-border 2s infinite;
        }
        .gradient-bg {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        }
        .glass-effect {
            background: rgba(255, 255, 255, 0.1);
            backdrop-filter: blur(10px);
            border-radius: 10px;
            border: 1px solid rgba(255, 255, 255, 0.2);
        }
        .safe-gradient {
            background: linear-gradient(135deg, #84fab0 0%, #8fd3f4 100%);
        }
        .unsafe-gradient {
            background: linear-gradient(135deg, #fa709a 0%, #fee140 100%);
        }
    </style>
</head>
<body class="min-h-screen gradient-bg text-white">
    <div class="container mx-auto px-4 py-8">
        <header class="text-center mb-12">
            <div class="inline-block p-4 rounded-full glass-effect float-animation mb-6">
                <i class="fas fa-shield-alt text-5xl text-white"></i>
            </div>
            <h1 class="text-4xl md:text-5xl font-bold mb-4">AI Content Moderator</h1>
            <p class="text-xl text-gray-200 max-w-2xl mx-auto">
                Advanced, multilingual content classification tool powered by AI
            </p>
        </header>

        <main class="max-w-6xl mx-auto">
            <div class="grid grid-cols-1 lg:grid-cols-3 gap-8 mb-12">
                <!-- API Key Section -->
                <div class="lg:col-span-1">
                    <div class="glass-effect p-6 rounded-xl h-full">
                        <h2 class="text-2xl font-bold mb-4 flex items-center">
                            <i class="fas fa-key mr-2"></i> API Configuration
                        </h2>
                        <div class="mb-4">
                            <label class="block text-sm font-medium mb-2">API Key</label>
                            <div class="relative">
                                <input type="password" id="apiKey" placeholder="Enter your API key" 
                                    class="w-full px-4 py-3 rounded-lg bg-white/10 border border-white/20 focus:outline-none focus:ring-2 focus:ring-indigo-400 text-white">
                                <button id="toggleApiKey" class="absolute right-3 top-3 text-gray-300 hover:text-white">
                                    <i class="fas fa-eye"></i>
                                </button>
                            </div>
                        </div>
                        <div class="mb-4">
                            <label class="block text-sm font-medium mb-2">Model</label>
                            <select id="modelSelect" class="w-full px-4 py-3 rounded-lg bg-white/10 border border-white/20 focus:outline-none focus:ring-2 focus:ring-indigo-400 text-white">
                                <option value="gemma-3-270m-it" selected>Gemma 3 270M IT</option>
                            </select>
                        </div>
                        <div class="mt-6">
                            <h3 class="text-lg font-semibold mb-2">API Endpoints</h3>
                            <div class="bg-black/20 p-4 rounded-lg text-sm font-mono">
                                <div class="mb-2">POST /v1/moderations</div>
                                <div>POST /v1/moderations/batch</div>
                            </div>
                        </div>
                    </div>
                </div>

                <!-- Main Content Section -->
                <div class="lg:col-span-2">
                    <div class="glass-effect p-6 rounded-xl">
                        <h2 class="text-2xl font-bold mb-4 flex items-center">
                            <i class="fas fa-check-circle mr-2"></i> Content Analysis
                        </h2>
                        
                        <!-- Tabs -->
                        <div class="flex border-b border-white/20 mb-6">
                            <button id="singleTab" class="px-4 py-2 font-medium border-b-2 border-indigo-400 text-indigo-300 tab-active">
                                Single Text
                            </button>
                            <button id="batchTab" class="px-4 py-2 font-medium border-b-2 border-transparent text-gray-300 hover:text-white">
                                Batch Processing
                            </button>
                        </div>
                        
                        <!-- Single Text Tab -->
                        <div id="singleContent" class="tab-content">
                            <div class="mb-6">
                                <label class="block text-sm font-medium mb-2">Text to Analyze</label>
                                <textarea id="textInput" rows="6" placeholder="Enter any text in any language for content moderation analysis..."
                                    class="w-full px-4 py-3 rounded-lg bg-white/10 border border-white/20 focus:outline-none focus:ring-2 focus:ring-indigo-400 text-white resize-none"></textarea>
                            </div>
                            
                            <div class="grid grid-cols-1 md:grid-cols-3 gap-4 mb-6">
                                <div>
                                    <label class="block text-sm font-medium mb-2">Max Tokens</label>
                                    <input type="range" id="maxTokens" min="1" max="10" value="3" step="1"
                                        class="w-full accent-indigo-500">
                                    <div class="text-center mt-1">
                                        <span id="maxTokensValue" class="text-sm">3</span>
                                    </div>
                                </div>
                                <div>
                                    <label class="block text-sm font-medium mb-2">Temperature</label>
                                    <input type="range" id="temperature" min="0" max="1" value="0.1" step="0.1"
                                        class="w-full accent-indigo-500">
                                    <div class="text-center mt-1">
                                        <span id="temperatureValue" class="text-sm">0.1</span>
                                    </div>
                                </div>
                                <div>
                                    <label class="block text-sm font-medium mb-2">Top-p</label>
                                    <input type="range" id="topP" min="0.1" max="1" value="0.95" step="0.05"
                                        class="w-full accent-indigo-500">
                                    <div class="text-center mt-1">
                                        <span id="topPValue" class="text-sm">0.95</span>
                                    </div>
                                </div>
                            </div>
                            
                            <div class="flex space-x-4">
                                <button id="analyzeBtn" class="flex-1 bg-indigo-600 hover:bg-indigo-700 text-white font-bold py-3 px-6 rounded-lg transition duration-300 transform hover:scale-105 pulse-border">
                                    <i class="fas fa-search mr-2"></i> Analyze Content
                                </button>
                                <button id="clearBtn" class="bg-gray-600 hover:bg-gray-700 text-white font-bold py-3 px-6 rounded-lg transition duration-300">
                                    <i class="fas fa-trash mr-2"></i> Clear
                                </button>
                            </div>
                        </div>
                        
                        <!-- Batch Processing Tab -->
                        <div id="batchContent" class="tab-content hidden">
                            <div class="mb-6">
                                <label class="block text-sm font-medium mb-2">Texts to Analyze (one per line)</label>
                                <textarea id="batchTextInput" rows="8" placeholder="Enter multiple texts, one per line..."
                                    class="w-full px-4 py-3 rounded-lg bg-white/10 border border-white/20 focus:outline-none focus:ring-2 focus:ring-indigo-400 text-white resize-none"></textarea>
                            </div>
                            
                            <div class="flex space-x-4">
                                <button id="batchAnalyzeBtn" class="flex-1 bg-indigo-600 hover:bg-indigo-700 text-white font-bold py-3 px-6 rounded-lg transition duration-300 transform hover:scale-105 pulse-border">
                                    <i class="fas fa-layer-group mr-2"></i> Analyze Batch
                                </button>
                                <button id="batchClearBtn" class="bg-gray-600 hover:bg-gray-700 text-white font-bold py-3 px-6 rounded-lg transition duration-300">
                                    <i class="fas fa-trash mr-2"></i> Clear
                                </button>
                            </div>
                        </div>
                        
                        <!-- Results Section -->
                        <div id="resultsSection" class="mt-8 hidden">
                            <h3 class="text-xl font-bold mb-4 flex items-center">
                                <i class="fas fa-chart-bar mr-2"></i> Analysis Results
                            </h3>
                            <div id="resultsContainer" class="space-y-4">
                                <!-- Results will be dynamically inserted here -->
                            </div>
                        </div>
                    </div>
                </div>
            </div>

            <!-- Examples Section -->
            <div class="glass-effect p-6 rounded-xl mb-12">
                <h2 class="text-2xl font-bold mb-4 flex items-center">
                    <i class="fas fa-lightbulb mr-2"></i> Example Prompts
                </h2>
                <div class="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-4">
                    <div class="example-card bg-white/10 p-4 rounded-lg cursor-pointer hover:bg-white/20 transition duration-300">
                        <p class="text-sm">"Hello, how are you today? I hope you're having a wonderful time!"</p>
                    </div>
                    <div class="example-card bg-white/10 p-4 rounded-lg cursor-pointer hover:bg-white/20 transition duration-300">
                        <p class="text-sm">"I hate you and I will find you and hurt you badly."</p>
                    </div>
                    <div class="example-card bg-white/10 p-4 rounded-lg cursor-pointer hover:bg-white/20 transition duration-300">
                        <p class="text-sm">"C'est une belle journée pour apprendre la programmation et l'intelligence artificielle."</p>
                    </div>
                    <div class="example-card bg-white/10 p-4 rounded-lg cursor-pointer hover:bg-white/20 transition duration-300">
                        <p class="text-sm">"I can't take this anymore. I want to end everything and disappear forever."</p>
                    </div>
                    <div class="example-card bg-white/10 p-4 rounded-lg cursor-pointer hover:bg-white/20 transition duration-300">
                        <p class="text-sm">"¡Hola! Me encanta aprender nuevos idiomas y conocer diferentes culturas."</p>
                    </div>
                    <div class="example-card bg-white/10 p-4 rounded-lg cursor-pointer hover:bg-white/20 transition duration-300">
                        <p class="text-sm">"You're absolutely worthless and nobody will ever love someone like you."</p>
                    </div>
                </div>
            </div>

            <!-- Information Section -->
            <div class="glass-effect p-6 rounded-xl">
                <h2 class="text-2xl font-bold mb-4 flex items-center">
                    <i class="fas fa-info-circle mr-2"></i> About This Tool
                </h2>
                <div class="grid grid-cols-1 md:grid-cols-3 gap-6">
                    <div class="text-center">
                        <div class="inline-block p-3 rounded-full bg-indigo-500/20 mb-3">
                            <i class="fas fa-globe text-2xl text-indigo-300"></i>
                        </div>
                        <h3 class="text-lg font-semibold mb-2">Multilingual</h3>
                        <p class="text-gray-300">Supports content analysis in multiple languages with high accuracy.</p>
                    </div>
                    <div class="text-center">
                        <div class="inline-block p-3 rounded-full bg-indigo-500/20 mb-3">
                            <i class="fas fa-bolt text-2xl text-indigo-300"></i>
                        </div>
                        <h3 class="text-lg font-semibold mb-2">Fast Processing</h3>
                        <p class="text-gray-300">Optimized model for quick content analysis with real-time results.</p>
                    </div>
                    <div class="text-center">
                        <div class="inline-block p-3 rounded-full bg-indigo-500/20 mb-3">
                            <i class="fas fa-shield-alt text-2xl text-indigo-300"></i>
                        </div>
                        <h3 class="text-lg font-semibold mb-2">Secure</h3>
                        <p class="text-gray-300">API key authentication ensures your requests remain secure and private.</p>
                    </div>
                </div>
            </div>
        </main>

        <footer class="mt-12 text-center text-gray-300">
            <p>© 2023 AI Content Moderator. All rights reserved.</p>
        </footer>
    </div>

    <!-- Loading Modal -->
    <div id="loadingModal" class="fixed inset-0 bg-black/70 flex items-center justify-center z-50 hidden">
        <div class="glass-effect p-8 rounded-xl max-w-md w-full mx-4 text-center">
            <div class="mb-4">
                <div class="inline-block animate-spin rounded-full h-12 w-12 border-t-2 border-b-2 border-indigo-500"></div>
            </div>
            <h3 class="text-xl font-bold mb-2">Analyzing Content</h3>
            <p class="text-gray-300">Please wait while we process your request...</p>
        </div>
    </div>

    <script>
        // DOM Elements
        const singleTab = document.getElementById('singleTab');
        const batchTab = document.getElementById('batchTab');
        const singleContent = document.getElementById('singleContent');
        const batchContent = document.getElementById('batchContent');
        const apiKeyInput = document.getElementById('apiKey');
        const toggleApiKeyBtn = document.getElementById('toggleApiKey');
        const textInput = document.getElementById('textInput');
        const batchTextInput = document.getElementById('batchTextInput');
        const maxTokensSlider = document.getElementById('maxTokens');
        const maxTokensValue = document.getElementById('maxTokensValue');
        const temperatureSlider = document.getElementById('temperature');
        const temperatureValue = document.getElementById('temperatureValue');
        const topPSlider = document.getElementById('topP');
        const topPValue = document.getElementById('topPValue');
        const analyzeBtn = document.getElementById('analyzeBtn');
        const batchAnalyzeBtn = document.getElementById('batchAnalyzeBtn');
        const clearBtn = document.getElementById('clearBtn');
        const batchClearBtn = document.getElementById('batchClearBtn');
        const resultsSection = document.getElementById('resultsSection');
        const resultsContainer = document.getElementById('resultsContainer');
        const loadingModal = document.getElementById('loadingModal');
        const exampleCards = document.querySelectorAll('.example-card');

        // Tab switching
        singleTab.addEventListener('click', () => {
            singleTab.classList.add('border-indigo-400', 'text-indigo-300');
            singleTab.classList.remove('border-transparent', 'text-gray-300');
            batchTab.classList.add('border-transparent', 'text-gray-300');
            batchTab.classList.remove('border-indigo-400', 'text-indigo-300');
            singleContent.classList.remove('hidden');
            batchContent.classList.add('hidden');
        });

        batchTab.addEventListener('click', () => {
            batchTab.classList.add('border-indigo-400', 'text-indigo-300');
            batchTab.classList.remove('border-transparent', 'text-gray-300');
            singleTab.classList.add('border-transparent', 'text-gray-300');
            singleTab.classList.remove('border-indigo-400', 'text-indigo-300');
            batchContent.classList.remove('hidden');
            singleContent.classList.add('hidden');
        });

        // Toggle API key visibility
        toggleApiKeyBtn.addEventListener('click', () => {
            const type = apiKeyInput.getAttribute('type') === 'password' ? 'text' : 'password';
            apiKeyInput.setAttribute('type', type);
            toggleApiKeyBtn.innerHTML = type === 'password' ? '<i class="fas fa-eye"></i>' : '<i class="fas fa-eye-slash"></i>';
        });

        // Slider value updates
        maxTokensSlider.addEventListener('input', () => {
            maxTokensValue.textContent = maxTokensSlider.value;
        });

        temperatureSlider.addEventListener('input', () => {
            temperatureValue.textContent = temperatureSlider.value;
        });

        topPSlider.addEventListener('input', () => {
            topPValue.textContent = topPSlider.value;
        });

        // Example cards
        exampleCards.forEach(card => {
            card.addEventListener('click', () => {
                textInput.value = card.querySelector('p').textContent;
            });
        });

        // Clear buttons
        clearBtn.addEventListener('click', () => {
            textInput.value = '';
            resultsSection.classList.add('hidden');
        });

        batchClearBtn.addEventListener('click', () => {
            batchTextInput.value = '';
            resultsSection.classList.add('hidden');
        });

        // Analyze button
        analyzeBtn.addEventListener('click', async () => {
            const text = textInput.value.trim();
            if (!text) {
                showNotification('Please enter text to analyze', 'error');
                return;
            }

            const apiKey = apiKeyInput.value.trim();
            if (!apiKey) {
                showNotification('Please enter your API key', 'error');
                return;
            }

            showLoading(true);

            try {
                const response = await fetch('/v1/moderations', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/json',
                        'Authorization': `Bearer ${apiKey}`
                    },
                    body: JSON.stringify({
                        input: text,
                        model: document.getElementById('modelSelect').value
                    })
                });

                if (!response.ok) {
                    const errorData = await response.json();
                    throw new Error(errorData.detail || 'An error occurred');
                }

                const data = await response.json();
                displayResults([data.results[0]]);
            } catch (error) {
                showNotification(`Error: ${error.message}`, 'error');
            } finally {
                showLoading(false);
            }
        });

        // Batch analyze button
        batchAnalyzeBtn.addEventListener('click', async () => {
            const texts = batchTextInput.value.trim().split('\\n').filter(text => text.trim());
            if (texts.length === 0) {
                showNotification('Please enter at least one text to analyze', 'error');
                return;
            }

            const apiKey = apiKeyInput.value.trim();
            if (!apiKey) {
                showNotification('Please enter your API key', 'error');
                return;
            }

            showLoading(true);

            try {
                const response = await fetch('/v1/moderations/batch', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/json',
                        'Authorization': `Bearer ${apiKey}`
                    },
                    body: JSON.stringify({
                        inputs: texts,
                        model: document.getElementById('modelSelect').value
                    })
                });

                if (!response.ok) {
                    const errorData = await response.json();
                    throw new Error(errorData.detail || 'An error occurred');
                }

                const data = await response.json();
                displayResults(data.results);
            } catch (error) {
                showNotification(`Error: ${error.message}`, 'error');
            } finally {
                showLoading(false);
            }
        });

        // Display results
        function displayResults(results) {
            resultsContainer.innerHTML = '';
            
            results.forEach((result, index) => {
                const isFlagged = result.flagged;
                const cardClass = isFlagged ? 'unsafe-gradient' : 'safe-gradient';
                const icon = isFlagged ? 'fas fa-exclamation-triangle' : 'fas fa-check-circle';
                const statusText = isFlagged ? 'UNSAFE' : 'SAFE';
                const statusDesc = isFlagged ? 
                    'Content may contain inappropriate or harmful material.' : 
                    'Content appears to be safe and appropriate.';

                const categories = Object.entries(result.categories)
                    .filter(([_, value]) => value)
                    .map(([key, _]) => key.replace('/', ' '))
                    .join(', ');

                const resultCard = document.createElement('div');
                resultCard.className = `p-6 rounded-xl text-white ${cardClass} shadow-lg`;
                resultCard.innerHTML = `
                    <div class="flex items-start">
                        <div class="mr-4 mt-1">
                            <i class="${icon} text-3xl"></i>
                        </div>
                        <div class="flex-1">
                            <div class="flex justify-between items-start mb-2">
                                <h3 class="text-xl font-bold">${statusText}</h3>
                                <span class="text-sm bg-black/20 px-2 py-1 rounded">Result ${index + 1}</span>
                            </div>
                            <p class="mb-4">${statusDesc}</p>
                            <div class="bg-black/20 p-4 rounded-lg mb-4">
                                <p class="text-sm font-mono break-words">${result.text}</p>
                            </div>
                            ${isFlagged ? `
                                <div class="mb-3">
                                    <h4 class="font-semibold mb-1">Flagged Categories:</h4>
                                    <p class="text-sm">${categories}</p>
                                </div>
                            ` : ''}
                            <div class="grid grid-cols-2 md:grid-cols-4 gap-2 text-xs">
                                ${Object.entries(result.category_scores).map(([category, score]) => `
                                    <div class="bg-black/20 p-2 rounded">
                                        <div class="font-medium">${category.replace('/', ' ')}</div>
                                        <div class="w-full bg-gray-700 rounded-full h-1.5 mt-1">
                                            <div class="bg-white h-1.5 rounded-full" style="width: ${score * 100}%"></div>
                                        </div>
                                        <div class="text-right mt-1">${(score * 100).toFixed(0)}%</div>
                                    </div>
                                `).join('')}
                            </div>
                        </div>
                    </div>
                `;
                
                resultsContainer.appendChild(resultCard);
            });
            
            resultsSection.classList.remove('hidden');
            resultsSection.scrollIntoView({ behavior: 'smooth' });
        }

        // Show/hide loading modal
        function showLoading(show) {
            if (show) {
                loadingModal.classList.remove('hidden');
            } else {
                loadingModal.classList.add('hidden');
            }
        }

        // Show notification
        function showNotification(message, type = 'info') {
            const notification = document.createElement('div');
            notification.className = `fixed top-4 right-4 p-4 rounded-lg shadow-lg z-50 ${
                type === 'error' ? 'bg-red-500' : 'bg-indigo-500'
            } text-white`;
            notification.innerHTML = `
                <div class="flex items-center">
                    <i class="fas ${type === 'error' ? 'fa-exclamation-circle' : 'fa-info-circle'} mr-2"></i>
                    <span>${message}</span>
                </div>
            `;
            
            document.body.appendChild(notification);
            
            setTimeout(() => {
                notification.style.opacity = '0';
                notification.style.transition = 'opacity 0.5s';
                setTimeout(() => {
                    document.body.removeChild(notification);
                }, 500);
            }, 3000);
        }
    </script>
</body>
</html>
""")

# Initialize the model
print("Initializing model...")
with torch.inference_mode():
    _ = model.generate(
        **tokenizer(["Hello"], return_tensors="pt").to(model.device),
        max_new_tokens=1, do_sample=False, use_cache=True
    )

print("🚀 Starting AI Content Moderator API...")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)