main.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import re
import csv
from typing import Dict, List, Optional, Union, Set
from pydantic import BaseModel, Field
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import RedirectResponse
from contextlib import asynccontextmanager

# Define the model architecture (unchanged from original code)
class ImprovedHybridCNNBiLSTM(nn.Module):
    def __init__(self, embedding_dim, max_length, num_classes, dropout_rate=0.5):
        super(ImprovedHybridCNNBiLSTM, self).__init__()

        # CNN part with multiple filter sizes
        self.filter_sizes = [2, 3, 4, 5]
        self.num_filters = 64

        # Conv layers with batch normalization
        self.convs = nn.ModuleList([
            nn.Sequential(
                nn.Conv1d(in_channels=embedding_dim,
                         out_channels=self.num_filters,
                         kernel_size=fs),
                nn.BatchNorm1d(self.num_filters),
                nn.ReLU(),
                nn.Dropout(dropout_rate/2)
            )
            for fs in self.filter_sizes
        ])

        # BiLSTM part
        self.lstm = nn.LSTM(embedding_dim,
                           hidden_size=64,
                           bidirectional=True,
                           batch_first=True,
                           dropout=dropout_rate if dropout_rate > 0 else 0)

        # Calculated size of concatenated features
        cnn_output_dim = self.num_filters * len(self.filter_sizes)
        lstm_output_dim = 64 * 2  # bidirectional = 2 * hidden_size

        # Attention mechanism for LSTM outputs
        self.attention = nn.Sequential(
            nn.Linear(64 * 2, 64),
            nn.Tanh(),
            nn.Linear(64, 1)
        )

        # Dense layers with residual connection
        self.fc1 = nn.Linear(cnn_output_dim + lstm_output_dim, 128)
        self.bn1 = nn.BatchNorm1d(128)
        self.fc2 = nn.Linear(128, 128)  # Same dimension for residual
        self.bn2 = nn.BatchNorm1d(128)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout_rate)
        self.fc3 = nn.Linear(128, num_classes)

        # Initialize weights
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def attention_net(self, lstm_output):
        # lstm_output shape: [batch_size, seq_len, hidden_size*2]
        attention_weights = self.attention(lstm_output).squeeze(-1)  # [batch_size, seq_len]
        soft_attention = F.softmax(attention_weights, dim=1)  # [batch_size, seq_len]
        context = torch.bmm(soft_attention.unsqueeze(1), lstm_output).squeeze(1)  # [batch_size, hidden_size*2]
        return context
 
    def forward(self, x, lengths):
        """
        Forward pass for the hybrid CNN-BiLSTM model
        
        Args:
            x: Input tensor of shape [batch_size, max_length, embedding_dim]
            lengths: Tensor of actual sequence lengths
            
        Returns:
            Tensor of logits with shape [batch_size, num_classes]
        """
        batch_size, seq_len, embed_dim = x.size()
        
        # CNN part: reshape for Conv1d which expects [batch, channels, length]
        x_conv = x.transpose(1, 2)  # [batch_size, embedding_dim, max_length]
        
        # Apply CNN layers and max pooling
        conv_outputs = []
        for conv in self.convs:
            # Apply convolution
            conv_out = conv(x_conv)  # [batch_size, num_filters, seq_len - filter_size + 1]
            # Apply max pooling over time
            pooled = F.max_pool1d(conv_out, conv_out.size(2))  # [batch_size, num_filters, 1]
            conv_outputs.append(pooled.squeeze(2))  # [batch_size, num_filters]
        
        # Concatenate all CNN outputs
        cnn_features = torch.cat(conv_outputs, dim=1)  # [batch_size, num_filters * len(filter_sizes)]
        
        # BiLSTM part
        packed_input = nn.utils.rnn.pack_padded_sequence(
            x, lengths.cpu(), batch_first=True, enforce_sorted=False
        )
        packed_output, (hidden, _) = self.lstm(packed_input)
        lstm_output, _ = nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True)
        
        # Apply attention
        lstm_features = self.attention_net(lstm_output)  # [batch_size, hidden_size*2]
        
        # Concatenate CNN and LSTM features
        combined_features = torch.cat([cnn_features, lstm_features], dim=1)
        
        # First dense layer
        x = self.fc1(combined_features)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.dropout(x)
        
        # Second dense layer with residual connection
        residual = x
        x = self.fc2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = x + residual  # Residual connection
        
        # Output layer
        logits = self.fc3(x)
        
        return logits

# Text preprocessing function
def clean_text(text):
    """
    Clean the text using regex
    """
    if isinstance(text, str):
        # Convert to lowercase
        text = text.lower()

        # Remove URLs
        text = re.sub(r'https?://\S+|www\.\S+', '', text)

        # Remove HTML tags
        text = re.sub(r'<.*?>', '', text)

        # Remove special characters but keep important punctuation
        text = re.sub(r'[^\w\s.,!?]', '', text)
        text = re.sub(r'[ \n]+', ' ', text)
        text = re.sub(r"[=%;]", "", text)

        # Remove extra whitespace
        text = re.sub(r'\s+', ' ', text).strip()

        return text
    else:
        return ""

# Create a class to handle text embedding
class TextEmbedder:
    def __init__(self, word_to_idx, word_embeddings, embedding_dim=100, max_length=100):
        self.word_to_idx = word_to_idx
        self.word_embeddings = word_embeddings
        self.embedding_dim = embedding_dim
        self.max_length = max_length
        
        # Add UNK token embedding as the average of all embeddings
        if self.word_embeddings is not None and len(self.word_embeddings) > 0:
            self.unk_embedding = np.mean(self.word_embeddings, axis=0)
        else:
            self.unk_embedding = np.zeros(self.embedding_dim)
            
    def embed_text(self, text):
        """
        Convert text to embedding tensor
        """
        # Clean the text
        cleaned_text = clean_text(text)
        
        # Tokenize
        tokens = cleaned_text.split()
        
        # Truncate if needed
        if len(tokens) > self.max_length:
            tokens = tokens[:self.max_length]
            
        # Get token length
        length = len(tokens)
        
        # Get embeddings for each word
        word_embeddings = []
        for word in tokens:
            # Get embedding for current word
            if word in self.word_to_idx:
                embedding = torch.tensor(self.word_embeddings[self.word_to_idx[word]], dtype=torch.float)
            else:
                embedding = torch.tensor(self.unk_embedding, dtype=torch.float)
            word_embeddings.append(embedding)
            
        # If the text is empty, add a zero vector
        if len(word_embeddings) == 0:
            word_embeddings.append(torch.zeros(self.embedding_dim))
            length = 1
            
        # Convert to tensor
        embeddings = torch.stack(word_embeddings)
        
        # Pad if needed
        if embeddings.size(0) < self.max_length:
            padding = torch.zeros(self.max_length - embeddings.size(0), self.embedding_dim)
            embeddings = torch.cat([embeddings, padding], dim=0)
            
        return embeddings.unsqueeze(0), torch.tensor([length])  # Add batch dimension

# NEW CLASS: Text Censor
class TextCensor:
    def __init__(self, profanity_file_path):
        self.profanity_words = self.load_profanity_list(profanity_file_path)
        print(f"Loaded {len(self.profanity_words)} profane words from {profanity_file_path}")
    
    def load_profanity_list(self, file_path) -> Set[str]:
        """
        Load profanity words from a CSV file
        """
        profanity_words = set()
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                for row in reader:
                    if row and len(row) > 0:
                        # Add the word after converting to lowercase
                        profanity_words.add(row[0].lower())
            return profanity_words
        except Exception as e:
            print(f"Warning: Could not load profanity list: {str(e)}")
            return set()
    
    def censor_text(self, text) -> str:
        """
        Censor profane words in the text by replacing them with asterisks
        """
        if not text or not self.profanity_words:
            return text
        
        words = re.findall(r'\b\w+\b', text.lower())
        censored_text = text
        
        for word in words:
            if word.lower() in self.profanity_words:
                # Replace the word with asterisks
                censored_word = '*' * len(word)
                # Use regex to replace the whole word with boundaries
                pattern = r'\b' + re.escape(word) + r'\b'
                censored_text = re.sub(pattern, censored_word, censored_text, flags=re.IGNORECASE)
                
        return censored_text

# Create a ToxicityPredictor class
class ToxicityPredictor:
    def __init__(self, model_path, word_embeddings_path=None, profanity_file_path=None):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.word_embeddings_path = word_embeddings_path
        
        # Load the model and embeddings
        self.load_model(model_path)
        
        # Initialize text censor if profanity file is provided
        self.text_censor = None
        if profanity_file_path:
            self.text_censor = TextCensor(profanity_file_path)
        
    def load_model(self, model_path):
        """
        Load model, embeddings, and other necessary components
        """
        try:
            checkpoint = torch.load(model_path, map_location=self.device)
            
            # Extract model parameters
            self.class_names = checkpoint.get('class_names', [
                'TOXICITY', 'SEVERE_TOXICITY', 'INSULT', 
                'PROFANITY', 'IDENTITY_ATTACK', 'THREAT', 'NOT_TOXIC'
            ])
            self.embedding_dim = checkpoint.get('embedding_dim', 100)
            self.custom_thresholds = checkpoint.get('custom_thresholds', {
                "TOXICITY": 0.5,
                "SEVERE_TOXICITY": 0.6,
                "INSULT": 0.4,
                "PROFANITY": 0.5,
                "IDENTITY_ATTACK": 0.4,
                "THREAT": 0.4,
                "NOT_TOXIC": 0.1
            })
            
            # Create and load model
            self.model = ImprovedHybridCNNBiLSTM(
                embedding_dim=self.embedding_dim,
                max_length=100,
                num_classes=len(self.class_names),
                dropout_rate=0.0  # No dropout during inference
            )
            self.model.load_state_dict(checkpoint['model_state_dict'])
            self.model.to(self.device)
            self.model.eval()
            
            # Check if we need to load embeddings from a separate file
            if 'word_to_idx' in checkpoint and 'word_embeddings' in checkpoint:
                self.word_to_idx = checkpoint['word_to_idx']
                self.word_embeddings = checkpoint['word_embeddings']
                print("Using embeddings from checkpoint")
            elif self.word_embeddings_path:
                # Load embeddings from file
                self.word_to_idx, self.word_embeddings = self.load_word_embeddings(self.word_embeddings_path)
                print(f"Loaded embeddings from {self.word_embeddings_path}")
            else:
                # If no embeddings in checkpoint and no path provided
                raise ValueError("Embeddings not found in model checkpoint and no embeddings file provided")
                
            # Create the text embedder
            self.text_embedder = TextEmbedder(
                self.word_to_idx,
                self.word_embeddings,
                embedding_dim=self.embedding_dim
            )
            
            print(f"Model loaded successfully. Available classes: {self.class_names}")
            
        except Exception as e:
            raise RuntimeError(f"Failed to load model: {str(e)}")

    def load_word_embeddings(self, embeddings_path):
        """
        Load word embeddings from FastText file
        """
        word_to_idx = {}
        word_embeddings = []
        
        print(f"Loading word embeddings from {embeddings_path}")
        
        try:
            with open(embeddings_path, 'r', encoding='utf-8') as f:
                # Skip the first line if it contains metadata (dimensions etc.)
                line = f.readline().strip()
                if len(line.split()) <= 3:  # Skip if it's a header line
                    pass
                else:  # It's actually the first vector
                    parts = line.split(' ')
                    word = parts[0]
                    vector = np.array([float(val) for val in parts[1:]])
                    word_to_idx[word] = 0
                    word_embeddings.append(vector)
                
                # Process remaining lines
                idx = len(word_embeddings)
                for line in f:
                    parts = line.strip().split(' ')
                    if len(parts) < self.embedding_dim + 1:
                        continue  # Skip malformed lines
                    
                    word = parts[0]
                    vector = np.array([float(val) for val in parts[1:]])
                    
                    word_to_idx[word] = idx
                    word_embeddings.append(vector)
                    idx += 1
                    
                    # Print progress periodically
                    if idx % 10000 == 0:
                        print(f"Loaded {idx} word vectors")
            
            print(f"Loaded {len(word_embeddings)} word vectors with dimension {self.embedding_dim}")
            return word_to_idx, np.array(word_embeddings)
            
        except Exception as e:
            raise RuntimeError(f"Failed to load word embeddings: {str(e)}")
            
    def predict(self, text):
        """
        Predict toxicity levels for input text
        """
        with torch.no_grad():
            # Embed the text
            embeddings, lengths = self.text_embedder.embed_text(text)
            embeddings = embeddings.to(self.device)
            lengths = lengths.to(self.device)
            
            # Get model prediction
            outputs = self.model(embeddings, lengths)
            probabilities = torch.sigmoid(outputs).cpu().numpy()[0]
            
            # Apply thresholds
            predictions = {}
            for i, class_name in enumerate(self.class_names):
                threshold = self.custom_thresholds.get(class_name, 0.5)
                predictions[class_name] = {
                    "probability": float(probabilities[i]),
                    "is_detected": bool(probabilities[i] >= threshold)
                }
                
            return predictions
    
    def censor_text(self, text):
        """
        Censor profane words in the text
        """
        if self.text_censor is None:
            return text
        
        return self.text_censor.censor_text(text)

# Define Pydantic models for request/response
class TextRequest(BaseModel):
    text: str = Field(..., description="Text to analyze for toxicity", example="This is a test message")
    min_probability: Optional[float] = Field(0.0, description="Minimum probability to include in results (0-1)", ge=0.0, le=1.0)
    apply_censoring: Optional[bool] = Field(True, description="Whether to apply profanity censoring")

class ToxicityCategory(BaseModel):
    probability: float = Field(..., description="Probability score (0-1)")
    is_detected: bool = Field(..., description="Whether toxicity is detected based on threshold")

class ToxicityResponse(BaseModel):
    results: Dict[str, ToxicityCategory]
    summary: Dict[str, Union[bool, List[str]]] = Field(
        ..., 
        description="Summary of toxicity detection results"
    )
    censored_text: Optional[str] = Field(None, description="Text with profanity censored")

# Global predictor instance
predictor = None

# Define the lifespan context manager for app startup/shutdown
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load model at startup
    global predictor
    try:
        # Initialize the predictor with model, word embeddings, and profanity list
        predictor = ToxicityPredictor(
            "model/toxicity_model.pt", 
            "model/fasttext_vectors_dim_100.txt",
            "merge-profanity.csv"  # Added profanity file path
        )
        print("Model, embeddings, and profanity list loaded successfully!")
    except Exception as e:
        print(f"Error loading model, embeddings, or profanity list: {e}")
        # Service will start, but API calls will fail until model is loaded
    
    yield  # This is where FastAPI runs and serves requests
    
    # Cleanup (if needed) when the app is shutting down
    print("Shutting down application")

# Initialize FastAPI with lifespan handler
app = FastAPI(
    title="Toxicity Detection API with Censoring",
    description="API for detecting toxic content in text and censoring profanity",
    version="1.0.0",
    lifespan=lifespan
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods
    allow_headers=["*"],  # Allows all headers
)

@app.get("/")
async def root():
    return RedirectResponse(url="/docs")

@app.post("/predict", response_model=ToxicityResponse)
async def predict_toxicity(request: TextRequest):
    global predictor
    
    if predictor is None:
        raise HTTPException(status_code=503, detail="Model not loaded. Please try again later.")
    
    if not request.text or len(request.text.strip()) == 0:
        raise HTTPException(status_code=400, detail="Text cannot be empty")
    
    try:
        # Get predictions
        predictions = predictor.predict(request.text)
        
        # Filter results if min_probability is specified
        if request.min_probability > 0:
            filtered_results = {
                k: v for k, v in predictions.items() 
                if v["probability"] >= request.min_probability
            }
        else:
            filtered_results = predictions
        
        # Generate summary
        detected_categories = [
            category for category, data in filtered_results.items() 
            if data["is_detected"] and category != "NOT_TOXIC"
        ]
        
        is_toxic = len(detected_categories) > 0
        
        summary = {
            "is_toxic": is_toxic,
            "detected_categories": detected_categories
        }
        
        # Apply censoring if requested
        censored_text = None
        if request.apply_censoring:
            censored_text = predictor.censor_text(request.text)
        
        return {
            "results": filtered_results,
            "summary": summary,
            "censored_text": censored_text
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")

@app.post("/censor")
async def censor_text(request: TextRequest):
    """Endpoint to just censor text without toxicity prediction"""
    global predictor
    
    if predictor is None or predictor.text_censor is None:
        raise HTTPException(status_code=503, detail="Censoring service not available. Please try again later.")
    
    if not request.text or len(request.text.strip()) == 0:
        raise HTTPException(status_code=400, detail="Text cannot be empty")
    
    try:
        censored_text = predictor.censor_text(request.text)
        return {
            "original_text": request.text,
            "censored_text": censored_text
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Censoring error: {str(e)}")

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    global predictor
    return {
        "status": "healthy",
        "model_loaded": predictor is not None,
        "censoring_available": predictor is not None and predictor.text_censor is not None
    }

@app.get("/info")
async def model_info():
    """Get information about the loaded model and censoring service"""
    global predictor
    
    if predictor is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    
    return {
        "model_type": "Hybrid CNN+BiLSTM",
        "classes": predictor.class_names,
        "custom_thresholds": predictor.custom_thresholds,
        "device": str(predictor.device),
        "embedding_source": "checkpoint" if predictor.word_embeddings_path is None else predictor.word_embeddings_path,
        "censoring_enabled": predictor.text_censor is not None
    }

# Run the API with uvicorn when the script is executed directly
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)