Update app.py

app.py

@@ -1,243 +1,30 @@
-from fastapi import FastAPI, HTTPException, Depends, Header
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from transformers import DistilBertForSequenceClassification, DistilBertTokenizer
-import torch
-from typing import Dict, List, Optional
-import os
-from firebase_service import firebase_service
+import gradio as gr
+from transformers import pipeline
 
-app = FastAPI(title="CyberGuard AI API", version="1.0.0")
+print("🔄 Loading CyberGuard Model...")
+classifier = pipeline("text-classification", model="Jishnuuuu/cyberguard-v1")
+print("✅ Model loaded successfully!")
 
-# Enable CORS for React Native app
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # In production, specify your app's origin
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Global variables for model and tokenizer
-model = None
-tokenizer = None
-device = None
-
-# Labels for classification (inverted based on model training)
-LABELS = {
-    0: "safe",
-    1: "threat"
-}
-
-SEVERITY_MAPPING = {
-    "safe": "none",
-    "threat": "high"
-}
-
-class MessageRequest(BaseModel):
-    text: str
-    app_name: Optional[str] = None
-    
-class MessageResponse(BaseModel):
-    text: str
-    is_harmful: bool
-    label: str
-    confidence: float
-    severity: str
-    details: Dict
-
-class AuthenticatedUser(BaseModel):
-    uid: str
-    email: str
-    email_verified: bool
-
-class UserStatsResponse(BaseModel):
-    total_alerts: int
-    high_severity: int
-    medium_severity: int
-    low_severity: int
-
-# Authentication dependency
-async def get_current_user(authorization: Optional[str] = Header(None)) -> Optional[AuthenticatedUser]:
-    """Get current authenticated user from Firebase ID token"""
-    if not authorization or not authorization.startswith('Bearer '):
-        return None
-    
-    token = authorization.split('Bearer ')[1]
-    user_data = firebase_service.verify_user_token(token)
-    
-    if user_data:
-        return AuthenticatedUser(
-            uid=user_data['uid'],
-            email=user_data['email'],
-            email_verified=user_data['email_verified']
-        )
-    return None
-
-@app.on_event("startup")
-async def load_model():
-    """Load the model and tokenizer on startup"""
-    global model, tokenizer, device
+def analyze_message(text):
+    """Analyze message for cyberbullying threats"""
+    if not text or len(text.strip()) == 0:
+        return {"Label": "ERROR", "Confidence": 0.0}
     
     try:
-        # Get the model path
-        model_path = os.path.join(
-            os.path.dirname(os.path.dirname(__file__)),
-            "CyberGuard Model V1",
-            "final_cyberbullying_model"
-        )
-        
-        print(f"Loading model from: {model_path}")
-        
-        # Check if model exists
-        if not os.path.exists(model_path):
-            print(f"⚠️  Model path not found: {model_path}")
-            print("⚠️  Backend will run in demo mode without AI model")
-            return
-        
-        # Set device
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"Using device: {device}")
-        
-        # Load tokenizer and model
-        tokenizer = DistilBertTokenizer.from_pretrained(model_path)
-        model = DistilBertForSequenceClassification.from_pretrained(model_path)
-        model.to(device)
-        model.eval()
-        
-        print("✅ Model loaded successfully!")
-        
-    except Exception as e:
-        print(f"⚠️  Model loading failed: {str(e)}")
-        print("⚠️  Backend will run in demo mode without AI model")
-        # Don't raise - allow backend to start without model
-        model = None
-        tokenizer = None
-
-@app.get("/")
-async def root():
-    """Health check endpoint"""
-    return {
-        "status": "online",
-        "message": "CyberGuard AI API is running",
-        "model_loaded": model is not None
-    }
-
-@app.post("/analyze", response_model=MessageResponse)
-async def analyze_message(request: MessageRequest, current_user: Optional[AuthenticatedUser] = Depends(get_current_user)):
-    """Analyze a message for cyberbullying/threats"""
-    
-    if model is None or tokenizer is None:
-        raise HTTPException(status_code=503, detail="Model not loaded")
-    
-    if not request.text or len(request.text.strip()) == 0:
-        raise HTTPException(status_code=400, detail="Text cannot be empty")
-    
-    try:
-        # Tokenize input
-        inputs = tokenizer(
-            request.text,
-            return_tensors="pt",
-            truncation=True,
-            max_length=512,
-            padding=True
-        )
-        
-        # Move to device
-        inputs = {k: v.to(device) for k, v in inputs.items()}
-        
-        # Get prediction
-        with torch.no_grad():
-            outputs = model(**inputs)
-            logits = outputs.logits
-            probabilities = torch.nn.functional.softmax(logits, dim=-1)
-            predicted_class = torch.argmax(probabilities, dim=-1).item()
-            confidence = probabilities[0][predicted_class].item()
-        
-        # Get label
-        label = LABELS.get(predicted_class, "unknown")
-        is_harmful = label == "threat"
-        severity = SEVERITY_MAPPING.get(label, "none")
-        
-        # Additional details
-        details = {
-            "safe_probability": round(probabilities[0][0].item(), 4),
-            "threat_probability": round(probabilities[0][1].item(), 4),
-            "model": "DistilBERT",
-            "version": "1.0"
+        result = classifier(text)[0]
+        return {
+            "Label": result['label'],
+            "Confidence": round(result['score'], 3)
         }
-        
-        # Store alert if user is authenticated and threat is detected
-        if current_user and is_harmful:
-            alert_data = {
-                'timestamp': None,  # Will be set by Firebase
-                'severity': severity,
-                'threat_type': 'cyberbullying',
-                'app_name': request.app_name,
-                'confidence': round(confidence, 4),
-                'model_version': '1.0'
-            }
-            firebase_service.store_threat_alert(current_user.uid, alert_data)
-        
-        return MessageResponse(
-            text=request.text,
-            is_harmful=is_harmful,
-            label=label,
-            confidence=round(confidence, 4),
-            severity=severity,
-            details=details
-        )
-        
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")
-
-@app.post("/analyze-batch")
-async def analyze_batch(messages: List[str]):
-    """Analyze multiple messages at once"""
-    
-    if model is None or tokenizer is None:
-        raise HTTPException(status_code=503, detail="Model not loaded")
-    
-    results = []
-    for text in messages:
-        if text and len(text.strip()) > 0:
-            response = await analyze_message(MessageRequest(text=text))
-            results.append(response.dict())
-    
-    return {"results": results, "count": len(results)}
-
-@app.get("/user/stats", response_model=UserStatsResponse)
-async def get_user_stats(current_user: AuthenticatedUser = Depends(get_current_user)):
-    """Get user's threat alert statistics"""
-    if not current_user:
-        raise HTTPException(status_code=401, detail="Authentication required")
-    
-    stats = firebase_service.get_user_alert_stats(current_user.uid)
-    return UserStatsResponse(**stats)
-
-@app.delete("/user/data")
-async def delete_user_data(current_user: AuthenticatedUser = Depends(get_current_user)):
-    """Delete all user data (GDPR compliance)"""
-    if not current_user:
-        raise HTTPException(status_code=401, detail="Authentication required")
-    
-    success = firebase_service.delete_user_data(current_user.uid)
-    if success:
-        return {"message": "User data deleted successfully"}
-    else:
-        raise HTTPException(status_code=500, detail="Failed to delete user data")
-
-@app.get("/health")
-async def health_check():
-    """Detailed health check"""
-    return {
-        "status": "healthy",
-        "model_loaded": model is not None,
-        "tokenizer_loaded": tokenizer is not None,
-        "device": str(device) if device else "unknown",
-        "firebase_initialized": firebase_service.is_initialized()
-    }
+        return {"Label": "ERROR", "Confidence": 0.0, "error": str(e)}
+
+demo = gr.Interface(
+    fn=analyze_message,
+    inputs=gr.Textbox(label="Message", placeholder="Type a message..."),
+    outputs=gr.JSON(label="Result"),
+    title="🛡️ CyberGuard AI",
+    description="Cyberbullying Detection - Model: Jishnuuuu/cyberguard-v1"
+)
 
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8001)
+demo.launch()