Update app.py

app.py

@@ -1,148 +1,92 @@
-from flask import Flask, request, jsonify, Response
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
-import torch.nn.functional as F
-import threading
-import time
-import queue
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse, JSONResponse
+from transformers import pipeline
 from nltk.tokenize import sent_tokenize
-import os
-import json
+import asyncio
+import uuid
 import logging
+import json
+from typing import Dict
+from collections import deque
 
-app = Flask(__name__)
+app = FastAPI()
 
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s %(levelname)s %(threadName)s %(message)s'
-)
+# Logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
 logger = logging.getLogger(__name__)
 
-# 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
-
-def load_model():
-    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)
-        probs = F.softmax(outputs.logits, dim=-1)
-        pred = torch.argmax(probs, dim=-1).item()
-        confidence = probs[0][pred].item()
-        return {
-            "text": text.strip(),
-            "label": labels[pred],
-            "confidence": round(confidence * 100, 2)
-        }
-
-def background_worker(user_id, text):
-    logger.info(f"Processing started for user_id={user_id}")
-    sessions[user_id]['status'] = "processing"
+# Load pipeline once at startup
+classifier = pipeline("text-classification", model="priyabrat/AI.or.Human.text.classification")
+sessions: Dict[str, Dict] = {}
+queues: Dict[str, deque] = {}
 
-    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)
-            logger.info(f"user_id={user_id} line={i} classified as {result['label']} ({result['confidence']}%)")
-            result["line"] = i
-            queues[user_id].put(f"data: {json.dumps(result)}\n\n")
-            time.sleep(0.1)
-
-        queues[user_id].put("event: done\ndata: Session complete\n\n")
-    except Exception as e:
-        logger.error(f"Error processing user_id={user_id}: {e}")
-        queues[user_id].put(f"event: error\ndata: {str(e)}\n\n")
-    finally:
-        sessions[user_id]['status'] = "done"
-        logger.info(f"Processing finished for user_id={user_id}")
-        time.sleep(1)
-        sessions.pop(user_id, None)
-        queues.pop(user_id, None)
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}
 
-@app.route('/start-session', methods=['POST'])
-def start_session():
-    data = request.get_json()
+@app.post("/start-session")
+async def start_session(request: Request, background_tasks: BackgroundTasks):
+    data = await request.json()
     user_id = data.get("user_id")
     text = data.get("text")
 
     if not user_id or not text:
-        logger.warning("Missing user_id or text in start-session request")
-        return jsonify({"error": "user_id and text are required"}), 400
+        return JSONResponse(content={"error": "user_id and text are required"}, status_code=400)
 
     if user_id in sessions:
-        logger.warning(f"Session already exists for user_id={user_id}")
-        return jsonify({"message": "Session already exists", "status": sessions[user_id]["status"]}), 409
+        return JSONResponse(content={"message": "Session already exists", "status": sessions[user_id]["status"]}, status_code=409)
 
-    logger.info(f"Starting session for user_id={user_id}")
-    sessions[user_id] = {"status": "pending"}
-    queues[user_id] = queue.Queue()
-    threading.Thread(target=background_worker, args=(user_id, text), daemon=True).start()
+    sessions[user_id] = {"status": "processing"}
+    queues[user_id] = deque()
 
-    return jsonify({"message": "Session started", "status": "pending"}), 202
+    background_tasks.add_task(process_text, user_id, text)
+    return {"message": "Session started", "status": "processing"}
+
+async def process_text(user_id: str, text: str):
+    try:
+        lines = sent_tokenize(text) if '\n' not in text else [l.strip() for l in text.strip().split('\n') if l.strip()]
+        for idx, line in enumerate(lines, 1):
+            result = classifier(line)[0]
+            label = result['label']
+            confidence = round(result['score'] * 100, 2)
+            payload = {
+                "line": idx,
+                "text": line,
+                "label": "AI-generated" if label == "LABEL_0" else "Human-written",
+                "confidence": confidence
+            }
+            queues[user_id].append(f"data: {json.dumps(payload)}\n\n")
+            await asyncio.sleep(0.1)
+
+        queues[user_id].append("event: done\ndata: Session complete\n\n")
+    except Exception as e:
+        logger.error(f"Error: {e}")
+        queues[user_id].append(f"event: error\ndata: {str(e)}\n\n")
+    finally:
+        sessions[user_id]['status'] = "done"
+        await asyncio.sleep(1)
+        sessions.pop(user_id, None)
+        queues.pop(user_id, None)
 
-@app.route('/stream/<user_id>')
-def stream(user_id):
+@app.get("/stream/{user_id}")
+async def stream(user_id: str):
     if user_id not in sessions:
-        logger.warning(f"No active session for user_id={user_id} in stream request")
-        return jsonify({"error": "No active session for this user"}), 404
+        return JSONResponse(content={"error": "No active session"}, status_code=404)
 
-    def event_stream():
+    async def event_stream():
         while True:
-            try:
-                message = queues[user_id].get(timeout=30)
-                yield message
-                if "event: done" in message or "event: error" in message:
-                    logger.info(f"Stream ended for user_id={user_id} with message: {message.strip()}")
-                    break
-            except queue.Empty:
-                logger.warning(f"Stream timeout for user_id={user_id}")
-                yield "event: timeout\ndata: No activity\n\n"
+            if queues[user_id]:
+                yield queues[user_id].popleft()
+            elif sessions[user_id]['status'] == 'done':
                 break
+            await asyncio.sleep(0.1)
 
-    return Response(event_stream(), mimetype="text/event-stream")
+    return StreamingResponse(event_stream(), media_type="text/event-stream")
 
-@app.route('/status/<user_id>')
-def session_status(user_id):
-    status = sessions.get(user_id, {}).get("status", "no_session")
-    logger.info(f"Status request for user_id={user_id}: {status}")
-    return jsonify({"status": status})
+@app.get("/status/{user_id}")
+async def status(user_id: str):
+    return {"status": sessions.get(user_id, {}).get("status", "no_session")}
 
-if __name__ == '__main__':
-    logger.info("Starting Flask app")
-    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 8080)))
+@app.get("/")
+async def index():
+    return {"message": "Server is running"}