Update app.py

app.py

@@ -19,33 +19,42 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Health check endpoint
+# Lazy-loaded shared state
+model = None
+tokenizer = None
+device = None
+labels = ["AI-generated", "Human-written"]
+lock = threading.Lock()
+
+sessions = {}
+queues = {}
+
+@app.route('/')
+def index():
+    logger.info("Index page requested")
+    return "Server is running!"
+
 @app.route('/health')
 def health_check():
     logger.info("Health check requested")
     return jsonify({"status": "healthy"}), 200
 
-# Initialize model only when needed
 def load_model():
-    model_name = "priyabrat/AI.or.Human.text.classification"
-    logger.info(f"Loading model and tokenizer from {model_name}")
-    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir='/app/hf_cache')
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, cache_dir='/app/hf_cache')
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device).eval()
-    logger.info(f"Model loaded on device: {device}")
-    return tokenizer, model, device
-
-
-tokenizer, model, device = load_model()
-labels = ["AI-generated", "Human-written"]
-lock = threading.Lock()
-
-sessions = {}
-queues = {}
+    global tokenizer, model, device
+    if model is None or tokenizer is None:
+        model_name = "priyabrat/AI.or.Human.text.classification"
+        logger.info(f"Loading model and tokenizer from {model_name}")
+        tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir='/app/hf_cache')
+        model = AutoModelForSequenceClassification.from_pretrained(model_name, cache_dir='/app/hf_cache')
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model.to(device).eval()
+        logger.info(f"Model loaded on device: {device}")
+    else:
+        logger.info("Model already loaded.")
 
 def classify_line(text):
     with lock, torch.no_grad():
+        load_model()
         inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
         inputs = {k: v.to(device) for k, v in inputs.items()}
         outputs = model(**inputs)
@@ -61,7 +70,7 @@ def classify_line(text):
 def background_worker(user_id, text):
     logger.info(f"Processing started for user_id={user_id}")
     sessions[user_id]['status'] = "processing"
-    
+
     try:
         if '\n' not in text:
             lines = sent_tokenize(text)
@@ -83,10 +92,8 @@ def background_worker(user_id, text):
         sessions[user_id]['status'] = "done"
         logger.info(f"Processing finished for user_id={user_id}")
         time.sleep(1)
-        if user_id in sessions:
-            del sessions[user_id]
-        if user_id in queues:
-            del queues[user_id]
+        sessions.pop(user_id, None)
+        queues.pop(user_id, None)
 
 @app.route('/start-session', methods=['POST'])
 def start_session():
@@ -136,11 +143,6 @@ def session_status(user_id):
     logger.info(f"Status request for user_id={user_id}: {status}")
     return jsonify({"status": status})
 
-@app.route('/')
-def index():
-    logger.info("Index page requested")
-    return "Server is running!"
-
 if __name__ == '__main__':
     logger.info("Starting Flask app")
     app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 8080)))