Update app.py

app.py

@@ -1,92 +1,83 @@
-from fastapi import FastAPI, Request, BackgroundTasks
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, StreamingResponse, PlainTextResponse
 from transformers import pipeline
 from nltk.tokenize import sent_tokenize
-import asyncio
-import uuid
-import logging
+import time
 import json
-from typing import Dict
-from collections import deque
+import threading
+import queue
+import logging
+import os
 
 app = FastAPI()
-
-# Logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Load pipeline once at startup
 classifier = pipeline("text-classification", model="priyabrat/AI.or.Human.text.classification")
-sessions: Dict[str, Dict] = {}
-queues: Dict[str, deque] = {}
+sessions = {}
+queues = {}
+
+@app.get("/")
+async def index():
+    return PlainTextResponse("✅ FastAPI server running on Hugging Face Spaces!")
 
 @app.get("/health")
-async def health():
+async def health_check():
     return {"status": "healthy"}
 
 @app.post("/start-session")
-async def start_session(request: Request, background_tasks: BackgroundTasks):
+async def start_session(request: Request):
     data = await request.json()
     user_id = data.get("user_id")
     text = data.get("text")
 
     if not user_id or not text:
-        return JSONResponse(content={"error": "user_id and text are required"}, status_code=400)
+        return JSONResponse({"error": "user_id and text required"}, status_code=400)
 
     if user_id in sessions:
-        return JSONResponse(content={"message": "Session already exists", "status": sessions[user_id]["status"]}, status_code=409)
-
-    sessions[user_id] = {"status": "processing"}
-    queues[user_id] = deque()
-
-    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)
+        return JSONResponse({"message": "Session exists", "status": sessions[user_id]["status"]}, status_code=409)
+
+    sessions[user_id] = {"status": "pending"}
+    queues[user_id] = queue.Queue()
+
+    def worker():
+        try:
+            sessions[user_id]["status"] = "processing"
+            lines = sent_tokenize(text) if '\n' not in text else [l.strip() for l in text.split('\n') if l.strip()]
+            for i, line in enumerate(lines, 1):
+                result = classifier(line)[0]
+                queues[user_id].put(f"data: {json.dumps({'line': i, 'text': line, 'label': result['label'], 'confidence': round(result['score']*100,2)})}\n\n")
+                time.sleep(0.1)
+            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)
+            del sessions[user_id]
+            del queues[user_id]
+
+    threading.Thread(target=worker, daemon=True).start()
+    return {"message": "Session started", "status": "pending"}
 
 @app.get("/stream/{user_id}")
 async def stream(user_id: str):
     if user_id not in sessions:
-        return JSONResponse(content={"error": "No active session"}, status_code=404)
+        return JSONResponse({"error": "No active session"}, status_code=404)
 
-    async def event_stream():
+    def event_stream():
         while True:
-            if queues[user_id]:
-                yield queues[user_id].popleft()
-            elif sessions[user_id]['status'] == 'done':
+            try:
+                msg = queues[user_id].get(timeout=30)
+                yield msg
+                if "event: done" in msg or "event: error" in msg:
+                    break
+            except queue.Empty:
+                yield "event: timeout\ndata: No activity\n\n"
                 break
-            await asyncio.sleep(0.1)
 
     return StreamingResponse(event_stream(), media_type="text/event-stream")
 
 @app.get("/status/{user_id}")
-async def status(user_id: str):
+async def session_status(user_id: str):
     return {"status": sessions.get(user_id, {}).get("status", "no_session")}
-
-@app.get("/")
-async def index():
-    return {"message": "Server is running"}