Update app.py

app.py

@@ -4,6 +4,7 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import re
 import pandas as pd
 import gradio as gr
+import os
 
 # -----------------------------
 # MODEL INITIALIZATION
@@ -11,21 +12,30 @@ import gradio as gr
 MODEL_NAME = "desklib/ai-text-detector-v1.01"
 tokenizer = None
 model = None
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Force CPU if CUDA is not properly initialized to prevent crash
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    dtype = torch.float16  # Half precision for GPU speed/memory
+else:
+    device = torch.device("cpu")
+    dtype = torch.float32  # Full precision for CPU stability
 
 def get_model():
     global tokenizer, model
     if model is None:
         print(f"Loading Specialized Model: {MODEL_NAME} on {device}")
+        # Added low_cpu_mem_usage to prevent Build Exit Code 1 (OOM)
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
         model = AutoModelForSequenceClassification.from_pretrained(
             MODEL_NAME, 
             num_labels=1, 
-            ignore_mismatched_sizes=True
+            ignore_mismatched_sizes=True,
+            low_cpu_mem_usage=True,
+            torch_dtype=dtype
         ).to(device).eval()
     return tokenizer, model
 
-# UPDATED: Threshold set to 59% for visual triggers
 THRESHOLD = 0.59 
 
 # -----------------------------
@@ -70,9 +80,8 @@ def analyze(text):
     text = text.strip()
     if not text:
         return "—", "—", "<em>Please enter text...</em>", None
-        
+            
     word_count = len(text.split())
-    # Word count requirement restored to 300
     if word_count < 300:
         warning_msg = f"⚠️ <b>Insufficient Text:</b> Your input has {word_count} words. Please enter at least 300 words for an accurate analysis."
         return "Too Short", "N/A", f"<div style='color: #b80d0d; padding: 20px; border: 1px solid #b80d0d; border-radius: 8px;'>{warning_msg}</div>", None
@@ -104,7 +113,7 @@ def analyze(text):
     weighted_avg = sum(p * l for p, l in zip(probs, lengths)) / total_words if total_words > 0 else 0
 
     # -----------------------------
-    # HTML RECONSTRUCTION (Highlights trigger at THRESHOLD)
+    # HTML RECONSTRUCTION
     # -----------------------------
     highlighted_html = "<div style='font-family: sans-serif; line-height: 1.8;'>"
     prob_map = {idx: probs[i] for i, idx in enumerate(pure_sents_indices)}
@@ -113,10 +122,9 @@ def analyze(text):
         if block.startswith("\n") or block.isspace():
             highlighted_html += block.replace("\n", "<br>")
             continue
-        
+            
         if i in prob_map:
             score = prob_map[i]
-            # Color determined by the new 59% threshold
             if score >= THRESHOLD:
                 color, bg = "#b80d0d", "rgba(184, 13, 13, 0.15)" # RED
             else:
@@ -129,8 +137,8 @@ def analyze(text):
             )
         else:
             highlighted_html += block
-    highlighted_html += "</div>"
 
+    highlighted_html += "</div>"
     label = f"{weighted_avg:.1%} AI Probability"
     display_score = f"{weighted_avg:.2%}"
     
@@ -142,7 +150,7 @@ def analyze(text):
 # -----------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("## 🕵️ AI Detector Pro: Raw Mode")
-    gr.Markdown(f"Direct model output from **{MODEL_NAME}**. Visual highlight now triggers at **{THRESHOLD*100:.0f}%**.")
+    gr.Markdown(f"Visual highlight triggers at **{THRESHOLD*100:.0f}%**.")
     
     with gr.Row():
         with gr.Column(scale=3):
@@ -151,7 +159,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=1):
             verdict_out = gr.Label(label="Model Verdict (Raw)")
             score_out = gr.Label(label="Exact Weighted Probability")
-    
+            
     with gr.Tabs():
         with gr.TabItem("Visual Heatmap"):
             html_out = gr.HTML()