Update app.py

app.py

@@ -1,9 +1,10 @@
-import torch
-from transformers import GPT2LMHeadModel, GPT2TokenizerFast, GPT2Config
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends, Security
+from fastapi.security import HTTPBearer
 from pydantic import BaseModel
-from contextlib import asynccontextmanager
+from transformers import GPT2LMHeadModel, GPT2TokenizerFast, GPT2Config
+import torch
 import asyncio
+from contextlib import asynccontextmanager
 
 # FastAPI app instance
 app = FastAPI()
@@ -11,6 +12,9 @@ app = FastAPI()
 # Global model and tokenizer variables
 model, tokenizer = None, None
 
+# HTTPBearer instance for security
+bearer_scheme = HTTPBearer()
+
 # Function to load model and tokenizer
 def load_model():
     model_path = "./Ai-Text-Detector/model"
@@ -31,12 +35,8 @@ def load_model():
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global model, tokenizer
-    try:
-        model, tokenizer = load_model()
-        yield
-    except Exception as e:
-        print(f"Startup error: {str(e)}")
-        raise RuntimeError(f"Failed to start application: {str(e)}")
+    model, tokenizer = load_model()
+    yield
 
 # Attach startup loader
 app = FastAPI(lifespan=lifespan)
@@ -47,42 +47,42 @@ class TextInput(BaseModel):
 
 # Sync text classification
 def classify_text(sentence: str):
-    try:
-        inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True)
-        input_ids = inputs["input_ids"]
-        attention_mask = inputs["attention_mask"]
-
-        with torch.no_grad():
-            outputs = model(input_ids, attention_mask=attention_mask, labels=input_ids)
-            loss = outputs.loss
-            perplexity = torch.exp(loss).item()
-
-        if perplexity < 60:
-            result = "AI-generated"
-        elif perplexity < 80:
-            result = "Probably AI-generated"
-        else:
-            result = "Human-written"
-
-        return result, perplexity
-    except Exception as e:
-        raise RuntimeError(f"Error during text classification: {str(e)}")
+    inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True)
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+
+    with torch.no_grad():
+        outputs = model(input_ids, attention_mask=attention_mask, labels=input_ids)
+        loss = outputs.loss
+        perplexity = torch.exp(loss).item()
+
+    if perplexity < 60:
+        result = "AI-generated"
+    elif perplexity < 80:
+        result = "Probably AI-generated"
+    else:
+        result = "Human-written"
+
+    return result, perplexity
 
-# POST route to analyze text
+# POST route to analyze text with Bearer token
 @app.post("/analyze")
-async def analyze_text(data: TextInput):
+async def analyze_text(data: TextInput, token: str = Depends(bearer_scheme)):
     user_input = data.text.strip()
     if not user_input:
         raise HTTPException(status_code=400, detail="Text cannot be empty")
 
-    try:
-        result, perplexity = await asyncio.to_thread(classify_text, user_input)
-        return {
-            "result": result,
-            "perplexity": round(perplexity, 2),
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")
+    # The token is automatically extracted from the Authorization header
+    # You can validate the token here if needed
+    print(f"Received Bearer Token: {token}")
+
+    # Run classification asynchronously to prevent blocking
+    result, perplexity = await asyncio.to_thread(classify_text, user_input)
+    
+    return {
+        "result": result,
+        "perplexity": round(perplexity, 2),
+    }
 
 # Health check route
 @app.get("/health")
@@ -96,4 +96,4 @@ def index():
         "message": "FastAPI API is up.",
         "try": "/docs to test the API.",
         "status": "OK"
-    }
+    }