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, Tuple
import logging
import os
import gradio as gr
from fastapi.middleware.cors import CORSMiddleware
from concurrent.futures import ThreadPoolExecutor
from functools import partial
import time
from datetime import datetime

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

MAX_LENGTH = 512
MODEL_NAME = "microsoft/deberta-v3-small"
WINDOW_SIZE = 6
WINDOW_OVERLAP = 2
CONFIDENCE_THRESHOLD = 0.65
BATCH_SIZE = 8
MAX_WORKERS = 4

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)
        
        self.executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)

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

    def create_windows(self, sentences: List[str], window_size: int, overlap: int) -> List[str]:
        if len(sentences) < window_size:
            return [" ".join(sentences)]

        windows = []
        stride = window_size - overlap
        for i in range(0, len(sentences) - window_size + 1, stride):
            window = sentences[i:i + window_size]
            windows.append(" ".join(window))
        return windows

    def create_centered_windows(self, sentences: List[str], window_size: int) -> Tuple[List[str], List[List[int]]]:
        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)

            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):
        if not torch.cuda.is_available():
            torch.set_num_threads(MAX_WORKERS)
            torch.set_num_interop_threads(MAX_WORKERS)
            
        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):
        logger.info("Initializing model and tokenizer...")
        
        from transformers import DebertaV2TokenizerFast
        
        self.tokenizer = DebertaV2TokenizerFast.from_pretrained(
            self.model_name,
            model_max_length=MAX_LENGTH,
            use_fast=True
        )
        
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.model_name,
            num_labels=2
        ).to(self.device)
            
        model_path = "model_20250209_184929_acc1.0000.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 quick_scan(self, text: str) -> Dict:
        if not text.strip():
            return {
                'prediction': 'unknown',
                'confidence': 0.0,
                'num_windows': 0
            }

        sentences = self.processor.split_into_sentences(text)
        windows = self.processor.create_windows(sentences, WINDOW_SIZE, WINDOW_OVERLAP)

        predictions = []
        
        for i in range(0, len(windows), BATCH_SIZE):
            batch_windows = windows[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 idx, window in enumerate(batch_windows):
                    prediction = {
                        'window': window,
                        'human_prob': probs[idx][1].item(),
                        'ai_prob': probs[idx][0].item(),
                        'prediction': 'human' if probs[idx][1] > probs[idx][0] else 'ai'
                    }
                    predictions.append(prediction)

            del inputs, outputs, probs
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        if not predictions:
            return {
                'prediction': 'unknown',
                'confidence': 0.0,
                'num_windows': 0
            }

        avg_human_prob = sum(p['human_prob'] for p in predictions) / len(predictions)
        avg_ai_prob = sum(p['ai_prob'] for p in predictions) / len(predictions)

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

    def detailed_scan(self, text: str) -> Dict:
        text = text.rstrip()
        
        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 {}

        windows, window_sentence_indices = self.processor.create_centered_windows(sentences, WINDOW_SIZE)

        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))}

        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):
                    center_idx = len(indices) // 2
                    center_weight = 0.7
                    edge_weight = 0.3 / (len(indices) - 1)

                    for pos, sent_idx in enumerate(indices):
                        weight = center_weight if pos == center_idx else edge_weight
                        sentence_appearances[sent_idx] += weight
                        sentence_scores[sent_idx]['human_prob'] += weight * probs[window_idx][1].item()
                        sentence_scores[sent_idx]['ai_prob'] += weight * probs[window_idx][0].item()

            del inputs, outputs, probs
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        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]

                if i > 0 and i < len(sentences) - 1:
                    prev_human = sentence_scores[i-1]['human_prob'] / sentence_appearances[i-1]
                    prev_ai = sentence_scores[i-1]['ai_prob'] / sentence_appearances[i-1]
                    next_human = sentence_scores[i+1]['human_prob'] / sentence_appearances[i+1]
                    next_ai = sentence_scores[i+1]['ai_prob'] / sentence_appearances[i+1]

                    current_pred = 'human' if human_prob > ai_prob else 'ai'
                    prev_pred = 'human' if prev_human > prev_ai else 'ai'
                    next_pred = 'human' if next_human > next_ai else 'ai'

                    if current_pred != prev_pred or current_pred != next_pred:
                        smooth_factor = 0.1
                        human_prob = (human_prob * (1 - smooth_factor) + 
                                    (prev_human + next_human) * smooth_factor / 2)
                        ai_prob = (ai_prob * (1 - smooth_factor) + 
                                (prev_ai + next_ai) * smooth_factor / 2)

                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)
                })

        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:
        html_parts = []
        
        for pred in sentence_predictions:
            sentence = pred['sentence']
            confidence = pred['confidence']
            
            if confidence >= CONFIDENCE_THRESHOLD:
                if pred['prediction'] == 'human':
                    color = "#90EE90"
                else:
                    color = "#FFB6C6"
            else:
                if pred['prediction'] == 'human':
                    color = "#E8F5E9"
                else:
                    color = "#FFEBEE"
                    
            html_parts.append(f'<span style="background-color: {color};">{sentence}</span>')
            
        return " ".join(html_parts)

    def aggregate_predictions(self, predictions: List[Dict]) -> Dict:
        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, mode: str, classifier: TextClassifier) -> tuple:
    start_time = time.time()
    
    word_count = len(text.split())
    
    original_mode = mode
    if word_count < 200 and mode == "detailed":
        mode = "quick"
    
    if mode == "quick":
        result = classifier.quick_scan(text)
        
        quick_analysis = f"""
        PREDICTION: {result['prediction'].upper()}
        Confidence: {result['confidence']*100:.1f}%
        Windows analyzed: {result['num_windows']}
        """
        
        if original_mode == "detailed":
            quick_analysis += f"\n\nNote: Switched to quick mode because text contains only {word_count} words. Minimum 200 words required for detailed analysis."
        
        execution_time = (time.time() - start_time) * 1000
        
        return (
            text,
            "Quick scan mode - no sentence-level analysis available",
            quick_analysis
        )
    else:
        analysis = classifier.detailed_scan(text)
        
        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)
        
        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']}
        """
        
        execution_time = (time.time() - start_time) * 1000
        
        return (
            analysis['highlighted_text'],
            "\n".join(detailed_analysis),
            overall_result
        )

classifier = TextClassifier()

demo = gr.Interface(
    fn=lambda text, mode: analyze_text(text, mode, classifier),
    inputs=[
        gr.Textbox(
            lines=8,
            placeholder="Enter text to analyze...",
            label="Input Text"
        ),
        gr.Radio(
            choices=["quick", "detailed"],
            value="quick",
            label="Analysis Mode",
            info="Quick mode for faster analysis, Detailed mode for sentence-level analysis"
        )
    ],
    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. Choose between quick scan and detailed sentence-level analysis. 200+ words suggested for accurate predictions.",
    api_name="predict",
    flagging_mode="never"
)

app = demo.app

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["*"],
)

if __name__ == "__main__":
    demo.queue()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )