|
|
from fastapi import FastAPI, Request |
|
|
from fastapi.responses import JSONResponse, StreamingResponse, PlainTextResponse |
|
|
from transformers import pipeline |
|
|
from nltk.tokenize import sent_tokenize |
|
|
import time |
|
|
import json |
|
|
import threading |
|
|
import queue |
|
|
import logging |
|
|
import os |
|
|
|
|
|
app = FastAPI() |
|
|
logging.basicConfig(level=logging.INFO) |
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
classifier = pipeline("text-classification", model="priyabrat/AI.or.Human.text.classification") |
|
|
sessions = {} |
|
|
queues = {} |
|
|
|
|
|
@app.get("/") |
|
|
async def index(): |
|
|
return PlainTextResponse("✅ FastAPI server running on Hugging Face Spaces!") |
|
|
|
|
|
@app.get("/health") |
|
|
async def health_check(): |
|
|
return {"status": "healthy"} |
|
|
|
|
|
@app.post("/start-session") |
|
|
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({"error": "user_id and text required"}, status_code=400) |
|
|
|
|
|
if user_id in sessions: |
|
|
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({"error": "No active session"}, status_code=404) |
|
|
|
|
|
def event_stream(): |
|
|
while True: |
|
|
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 |
|
|
|
|
|
return StreamingResponse(event_stream(), media_type="text/event-stream") |
|
|
|
|
|
@app.get("/status/{user_id}") |
|
|
async def session_status(user_id: str): |
|
|
return {"status": sessions.get(user_id, {}).get("status", "no_session")} |
|
|
|
