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 # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Constants MAX_LENGTH = 512 MODEL_NAME = "microsoft/deberta-v3-small" WINDOW_SIZE = 6 WINDOW_OVERLAP = 2 CONFIDENCE_THRESHOLD = 0.65 BATCH_SIZE = 8 # Reduced batch size for CPU MAX_WORKERS = 4 # Number of worker threads for processing 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) # Initialize thread pool for parallel processing 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]]]: """Create windows with better boundary handling""" windows = [] window_sentence_indices = [] for i in range(len(sentences)): # Calculate window boundaries centered on current sentence half_window = window_size // 2 start_idx = max(0, i - half_window) end_idx = min(len(sentences), i + half_window + 1) # Create the window 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): # Set thread configuration before any model loading or parallel work 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): """Initialize the model and tokenizer.""" 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: """Perform a quick scan using simple window analysis.""" 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 = [] # Process windows in smaller batches for CPU efficiency 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) # Clean up GPU memory if available 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: """Original prediction method with modified window handling""" # Clean up trailing whitespace text = text.rstrip() if not text.strip(): return { 'sentence_predictions': [], 'highlighted_text': '', 'full_text': '', 'overall_prediction': { 'prediction': 'unknown', 'confidence': 0.0, 'num_sentences': 0 } } self.model.eval() 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 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) # Attribute predictions with weighted scoring for window_idx, indices in enumerate(batch_indices): center_idx = len(indices) // 2 center_weight = 0.7 # Higher weight for center sentence edge_weight = 0.3 / (len(indices) - 1) # Distribute remaining weight for pos, sent_idx in enumerate(indices): # Apply higher weight to center sentence 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() # Clean up memory del inputs, outputs, probs if torch.cuda.is_available(): torch.cuda.empty_cache() # Calculate final predictions with boundary smoothing 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] # Only apply minimal smoothing at prediction boundaries 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] # Check if we're at a prediction boundary 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: # Small adjustment at boundaries 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: """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'{sentence}') 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, mode: str, classifier: TextClassifier) -> tuple: """Analyze text using specified mode and return formatted results.""" 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']} """ return ( text, # No highlighting in quick mode "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']} """ 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, 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.", 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.", "quick"], ["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.", "detailed"], ], api_name="predict", flagging_mode="never" ) app = demo.app app.add_middleware( CORSMiddleware, allow_origins=["*"], # For development - in production, specify exact domains allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Launch the interface if __name__ == "__main__": demo.queue() # Enable queuing demo.launch( server_name="0.0.0.0", # Allow external connections server_port=7860, share=False # Don't use share since you're on Spaces )