Update app.py

app.py

@@ -5,33 +5,35 @@ import torch.nn.functional as F
 import threading
 import time
 import queue
-from nltk.tokenize import sent_tokenize 
-# import nltk
-# try:
-#     nltk.data.find('tokenizers/punkt')
-# except LookupError:
-#     nltk.download('punkt')
-
+from nltk.tokenize import sent_tokenize
+import os
 
 app = Flask(__name__)
 
-
-model_name = "priyabrat/AI.or.Human.text.classification"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device).eval()
-
+# Health check endpoint
+@app.route('/health')
+def health_check():
+    return jsonify({"status": "healthy"}), 200
+
+# Initialize model only when needed
+def load_model():
+    model_name = "priyabrat/AI.or.Human.text.classification"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device).eval()
+    return tokenizer, model
+
+tokenizer, model = load_model()
 labels = ["AI-generated", "Human-written"]
 lock = threading.Lock()
 
-
 sessions = {}
 queues = {}
 
 def classify_line(text):
     with lock, torch.no_grad():
-        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=10000)
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)  # Reduced max_length
         inputs = {k: v.to(device) for k, v in inputs.items()}
         outputs = model(**inputs)
         probs = F.softmax(outputs.logits, dim=-1)
@@ -43,47 +45,31 @@ def classify_line(text):
             "confidence": round(confidence * 100, 2)
         }
 
-
-
 def background_worker(user_id, text):
     sessions[user_id]['status'] = "processing"
-    if '\n' not in text:
-        lines = sent_tokenize(text)  
-    else:
-        lines = [line.strip() for line in text.strip().split('\n') if line.strip()]
-
-    result_count = 0
-
-    for i, line in enumerate(lines, 1):
-        result = classify_line(line)
-        result["line"] = i
-        queues[user_id].put(f"data: {result}\n\n")
-        result_count += 1
-        time.sleep(0.2)  
-
-    queues[user_id].put("event: done\ndata: Session complete\n\n")
-    sessions[user_id]['status'] = "done"
-
-    time.sleep(2)
-    del sessions[user_id]
-    del queues[user_id]
-
-    sessions[user_id]['status'] = "processing"
-    lines = [line.strip() for line in text.strip().split('\n') if line.strip()]
-    result_count = 0
-
-    for i, line in enumerate(lines, 1):
-        result = classify_line(line)
-        result["line"] = i
-        queues[user_id].put(f"data: {result}\n\n")
-        result_count += 1
-        time.sleep(0.2)  
-
-    queues[user_id].put("event: done\ndata: Session complete\n\n")
-    sessions[user_id]['status'] = "done"
-    time.sleep(2)
-    del sessions[user_id]
-    del queues[user_id]
+    
+    try:
+        if '\n' not in text:
+            lines = sent_tokenize(text)  
+        else:
+            lines = [line.strip() for line in text.strip().split('\n') if line.strip()]
+
+        for i, line in enumerate(lines, 1):
+            result = classify_line(line)
+            result["line"] = i
+            queues[user_id].put(f"data: {json.dumps(result)}\n\n")
+            time.sleep(0.1)  # Reduced delay
+
+        queues[user_id].put("event: done\ndata: Session complete\n\n")
+    except Exception as e:
+        queues[user_id].put(f"event: error\ndata: {str(e)}\n\n")
+    finally:
+        sessions[user_id]['status'] = "done"
+        time.sleep(1)
+        if user_id in sessions:
+            del sessions[user_id]
+        if user_id in queues:
+            del queues[user_id]
 
 @app.route('/start-session', methods=['POST'])
 def start_session():
@@ -95,8 +81,7 @@ def start_session():
         return jsonify({"error": "user_id and text are required"}), 400
 
     if user_id in sessions:
-        status = sessions[user_id]["status"]
-        return jsonify({"message": f"Session already exists", "status": status}), 409
+        return jsonify({"message": "Session already exists", "status": sessions[user_id]["status"]}), 409
 
     sessions[user_id] = {"status": "pending"}
     queues[user_id] = queue.Queue()
@@ -112,34 +97,23 @@ def stream(user_id):
     def event_stream():
         while True:
             try:
-                message = queues[user_id].get(timeout=60)
+                message = queues[user_id].get(timeout=30)  # Reduced timeout
                 yield message
-                if "event: done" in message:
+                if "event: done" in message or "event: error" in message:
                     break
             except queue.Empty:
                 yield "event: timeout\ndata: No activity\n\n"
                 break
 
-    return Response(
-        event_stream(),
-        mimetype="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
-            "Access-Control-Allow-Origin": "*"
-        }
-    )
+    return Response(event_stream(), mimetype="text/event-stream")
+
 @app.route('/status/<user_id>')
 def session_status(user_id):
-    if user_id not in sessions:
-        return jsonify({"status": "no_session"})
-    return jsonify({
-        "status": sessions[user_id]["status"]
-    })
+    return jsonify({"status": sessions.get(user_id, {}).get("status", "no_session")})
 
 @app.route('/')
 def index():
-    return "alive yet !"
+    return "Server is running!"
 
 if __name__ == '__main__':
-    app.run(threaded=True,host='0.0.0.0', port=8080)
+    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 8080)))