Update app.py

app.py

@@ -13,18 +13,16 @@ 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  
+BATCH_SIZE = 8
+MAX_WORKERS = 4
 
 class TextWindowProcessor:
     def __init__(self):
@@ -41,7 +39,6 @@ class TextWindowProcessor:
         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]:
@@ -64,12 +61,10 @@ class TextWindowProcessor:
         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)))
@@ -78,7 +73,6 @@ class TextWindowProcessor:
 
 class TextClassifier:
     def __init__(self):
-        
         if not torch.cuda.is_available():
             torch.set_num_threads(MAX_WORKERS)
             torch.set_num_interop_threads(MAX_WORKERS)
@@ -91,7 +85,6 @@ class TextClassifier:
         self.initialize_model()
 
     def initialize_model(self):
-        """Initialize the model and tokenizer."""
         logger.info("Initializing model and tokenizer...")
         
         from transformers import DebertaV2TokenizerFast
@@ -130,7 +123,6 @@ class TextClassifier:
 
         predictions = []
         
-        
         for i in range(0, len(windows), BATCH_SIZE):
             batch_windows = windows[i:i + BATCH_SIZE]
 
@@ -155,7 +147,6 @@ class TextClassifier:
                     }
                     predictions.append(prediction)
 
-            
             del inputs, outputs, probs
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
@@ -177,7 +168,6 @@ class TextClassifier:
         }
 
     def detailed_scan(self, text: str) -> Dict:
-        
         text = text.rstrip()
         
         if not text.strip():
@@ -196,14 +186,11 @@ class TextClassifier:
         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]
@@ -220,45 +207,38 @@ class TextClassifier:
                 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)  
+                    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)
@@ -289,14 +269,14 @@ class TextClassifier:
             
             if confidence >= CONFIDENCE_THRESHOLD:
                 if pred['prediction'] == 'human':
-                    color = "
+                    color = "#90EE90"
                 else:
-                    color = "
+                    color = "#FFB6C6"
             else:
                 if pred['prediction'] == 'human':
-                    color = "
+                    color = "#E8F5E9"
                 else:
-                    color = "
+                    color = "#FFEBEE"
                     
             html_parts.append(f'<span style="background-color: {color};">{sentence}</span>')
             
@@ -324,13 +304,10 @@ class TextClassifier:
         }
 
 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"
@@ -344,15 +321,13 @@ def analyze_text(text: str, mode: str, classifier: TextClassifier) -> tuple:
         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,  
+            text,
             "Quick scan mode - no sentence-level analysis available",
             quick_analysis
         )
@@ -374,7 +349,6 @@ def analyze_text(text: str, mode: str, classifier: TextClassifier) -> tuple:
         Number of sentences analyzed: {final_pred['num_sentences']}
         """
         
-        
         execution_time = (time.time() - start_time) * 1000
         
         return (
@@ -383,10 +357,8 @@ def analyze_text(text: str, mode: str, classifier: TextClassifier) -> tuple:
             overall_result
         )
 
-
 classifier = TextClassifier()
 
-
 demo = gr.Interface(
     fn=lambda text, mode: analyze_text(text, mode, classifier),
     inputs=[
@@ -413,12 +385,11 @@ demo = gr.Interface(
     flagging_mode="never"
 )
 
-
 app = demo.app
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["GET", "POST", "OPTIONS"],
     allow_headers=["*"],