app.py

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F
import spacy
from typing import List, Dict
import logging
import os
import gradio as gr

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Constants
MAX_LENGTH = 512
MODEL_NAME = "microsoft/deberta-v3-small"
WINDOW_SIZE = 17
WINDOW_OVERLAP = 2
CONFIDENCE_THRESHOLD = 0.65

class TextWindowProcessor:
    def __init__(self):
        try:
            self.nlp = spacy.load("en_core_web_sm")
        except OSError:
            logger.info("Downloading spacy model...")
            spacy.cli.download("en_core_web_sm")
            self.nlp = spacy.load("en_core_web_sm")

        if 'sentencizer' not in self.nlp.pipe_names:
            self.nlp.add_pipe('sentencizer')

        disabled_pipes = [pipe for pipe in self.nlp.pipe_names if pipe != 'sentencizer']
        self.nlp.disable_pipes(*disabled_pipes)

    def split_into_sentences(self, text: str) -> List[str]:
        doc = self.nlp(text)
        return [str(sent).strip() for sent in doc.sents]

    def create_centered_windows(self, sentences: List[str], window_size: int) -> tuple[List[str], List[List[int]]]:
        """Create windows centered around each sentence for detailed analysis."""
        windows = []
        window_sentence_indices = []

        for i in range(len(sentences)):
            half_window = window_size // 2
            start_idx = max(0, i - half_window)
            end_idx = min(len(sentences), i + half_window + 1)

            if start_idx == 0:
                end_idx = min(len(sentences), window_size)
            elif end_idx == len(sentences):
                start_idx = max(0, len(sentences) - window_size)

            window = sentences[start_idx:end_idx]
            windows.append(" ".join(window))
            window_sentence_indices.append(list(range(start_idx, end_idx)))

        return windows, window_sentence_indices

class TextClassifier:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model_name = MODEL_NAME
        self.tokenizer = None
        self.model = None
        self.processor = TextWindowProcessor()
        self.initialize_model()
        
    def initialize_model(self):
            """Initialize the model and tokenizer."""
            logger.info("Initializing model and tokenizer...")
            
            # Download and save tokenizer files locally
            local_tokenizer_path = "tokenizer"
            if not os.path.exists(local_tokenizer_path):
                AutoTokenizer.from_pretrained(self.model_name).save_pretrained(local_tokenizer_path)
            
            # Load from local files
            self.tokenizer = AutoTokenizer.from_pretrained(local_tokenizer_path)
            
            # First initialize the base model
            self.model = AutoModelForSequenceClassification.from_pretrained(
                self.model_name,
                num_labels=2
            ).to(self.device)
            
            # Look for model file in the same directory as the code
            model_path = "model.pt"  # Your model file should be uploaded as model.pt
            if os.path.exists(model_path):
                logger.info(f"Loading custom model from {model_path}")
                checkpoint = torch.load(model_path, map_location=self.device)
                self.model.load_state_dict(checkpoint['model_state_dict'])
            else:
                logger.warning("Custom model file not found. Using base model.")
                
            self.model.eval()

    def predict_with_sentence_scores(self, text: str) -> Dict:
        """Predict with sentence-level granularity using overlapping windows."""
        if not text.strip():
            return {
                'sentence_predictions': [],
                'highlighted_text': '',
                'full_text': '',
                'overall_prediction': {
                    'prediction': 'unknown',
                    'confidence': 0.0,
                    'num_sentences': 0
                }
            }

        sentences = self.processor.split_into_sentences(text)
        if not sentences:
            return {}

        # Create centered windows for each sentence
        windows, window_sentence_indices = self.processor.create_centered_windows(sentences, WINDOW_SIZE)

        # Track scores for each sentence
        sentence_appearances = {i: 0 for i in range(len(sentences))}
        sentence_scores = {i: {'human_prob': 0.0, 'ai_prob': 0.0} for i in range(len(sentences))}

        # Process windows in batches to save memory
        batch_size = 16
        for i in range(0, len(windows), batch_size):
            batch_windows = windows[i:i + batch_size]
            batch_indices = window_sentence_indices[i:i + batch_size]

            inputs = self.tokenizer(
                batch_windows,
                truncation=True,
                padding=True,
                max_length=MAX_LENGTH,
                return_tensors="pt"
            ).to(self.device)

            with torch.no_grad():
                outputs = self.model(**inputs)
                probs = F.softmax(outputs.logits, dim=-1)

                for window_idx, indices in enumerate(batch_indices):
                    for sent_idx in indices:
                        sentence_appearances[sent_idx] += 1
                        sentence_scores[sent_idx]['human_prob'] += probs[window_idx][1].item()
                        sentence_scores[sent_idx]['ai_prob'] += probs[window_idx][0].item()

        # Average the scores and create final sentence-level predictions
        sentence_predictions = []
        for i in range(len(sentences)):
            if sentence_appearances[i] > 0:
                human_prob = sentence_scores[i]['human_prob'] / sentence_appearances[i]
                ai_prob = sentence_scores[i]['ai_prob'] / sentence_appearances[i]
                sentence_predictions.append({
                    'sentence': sentences[i],
                    'human_prob': human_prob,
                    'ai_prob': ai_prob,
                    'prediction': 'human' if human_prob > ai_prob else 'ai',
                    'confidence': max(human_prob, ai_prob)
                })

        # Generate analysis outputs
        return {
            'sentence_predictions': sentence_predictions,
            'highlighted_text': self.format_predictions_html(sentence_predictions),
            'full_text': text,
            'overall_prediction': self.aggregate_predictions(sentence_predictions)
        }

    def format_predictions_html(self, sentence_predictions: List[Dict]) -> str:
        """Format predictions as HTML with color-coding."""
        html_parts = []
        
        for pred in sentence_predictions:
            sentence = pred['sentence']
            confidence = pred['confidence']
            
            if confidence >= CONFIDENCE_THRESHOLD:
                if pred['prediction'] == 'human':
                    color = "#90EE90"  # Light green
                else:
                    color = "#FFB6C6"  # Light red
            else:
                if pred['prediction'] == 'human':
                    color = "#E8F5E9"  # Very light green
                else:
                    color = "#FFEBEE"  # Very light red
                    
            html_parts.append(f'<span style="background-color: {color};">{sentence}</span>')
            
        return " ".join(html_parts)

    def aggregate_predictions(self, predictions: List[Dict]) -> Dict:
        """Aggregate predictions from multiple sentences into a single prediction."""
        if not predictions:
            return {
                'prediction': 'unknown',
                'confidence': 0.0,
                'num_sentences': 0
            }

        total_human_prob = sum(p['human_prob'] for p in predictions)
        total_ai_prob = sum(p['ai_prob'] for p in predictions)
        num_sentences = len(predictions)

        avg_human_prob = total_human_prob / num_sentences
        avg_ai_prob = total_ai_prob / num_sentences

        return {
            'prediction': 'human' if avg_human_prob > avg_ai_prob else 'ai',
            'confidence': max(avg_human_prob, avg_ai_prob),
            'num_sentences': num_sentences
        }

def analyze_text(text: str, classifier: TextClassifier) -> tuple:
    """Analyze text and return formatted results for Gradio interface."""
    # Get predictions
    analysis = classifier.predict_with_sentence_scores(text)
    
    # Format sentence-by-sentence analysis
    detailed_analysis = []
    for pred in analysis['sentence_predictions']:
        confidence = pred['confidence'] * 100
        detailed_analysis.append(f"Sentence: {pred['sentence']}")
        detailed_analysis.append(f"Prediction: {pred['prediction'].upper()}")
        detailed_analysis.append(f"Confidence: {confidence:.1f}%")
        detailed_analysis.append("-" * 50)
    
    # Format overall prediction
    final_pred = analysis['overall_prediction']
    overall_result = f"""
    FINAL PREDICTION: {final_pred['prediction'].upper()}
    Overall confidence: {final_pred['confidence']*100:.1f}%
    Number of sentences analyzed: {final_pred['num_sentences']}
    """
    
    return (
        analysis['highlighted_text'],
        "\n".join(detailed_analysis),
        overall_result
    )

# Initialize the classifier globally
classifier = TextClassifier()

# Create Gradio interface
demo = gr.Interface(
    fn=lambda text: analyze_text(text, classifier),
    inputs=gr.Textbox(
        lines=8,
        placeholder="Enter text to analyze...",
        label="Input Text"
    ),
    outputs=[
        gr.HTML(label="Highlighted Analysis"),
        gr.Textbox(label="Sentence-by-Sentence Analysis", lines=10),
        gr.Textbox(label="Overall Result", lines=4)
    ],
    title="AI Text Detector",
    description="Analyze text to detect if it was written by a human or AI. Text is analyzed sentence by sentence, with color coding indicating the prediction confidence.",
    examples=[
        ["This is a sample text written by a human. It contains multiple sentences with different ideas. The analysis will show how each sentence is classified. This demonstrates the AI detection capabilities."],
    ],
    allow_flagging="never"
)

# Launch the interface
if __name__ == "__main__":
    demo.launch()