Merge pull request #4 from cyberalertnepal/Pujan

Fix server crashes and high perplexity issue, refactor app.py, and update documentation

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
+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,9 +12,10 @@ app = FastAPI()
 # Global model and tokenizer variables
 model, tokenizer = None, None
 
-# Function to load model and tokenizer
-
+# HTTPBearer instance for security
+bearer_scheme = HTTPBearer()
 
+# Function to load model and tokenizer
 def load_model():
     model_path = "./Ai-Text-Detector/model"
     weights_path = "./Ai-Text-Detector/model_weights.pth"
@@ -22,39 +24,28 @@ def load_model():
         tokenizer = GPT2TokenizerFast.from_pretrained(model_path)
         config = GPT2Config.from_pretrained(model_path)
         model = GPT2LMHeadModel(config)
-        model.load_state_dict(
-            torch.load(weights_path, map_location=torch.device("cpu"))
-        )
-        model.eval()  # Set model to evaluation mode
+        model.load_state_dict(torch.load(weights_path, map_location=torch.device("cpu")))
+        model.eval() 
     except Exception as e:
         raise RuntimeError(f"Error loading model: {str(e)}")
 
     return model, tokenizer
 
-
 # Load model on app startup
-
-
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global model, tokenizer
     model, tokenizer = load_model()
     yield
 
-
 # Attach startup loader
 app = FastAPI(lifespan=lifespan)
 
 # Input schema
-
-
 class TextInput(BaseModel):
     text: str
 
-
 # Sync text classification
-
-
 def classify_text(sentence: str):
     inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding=True)
     input_ids = inputs["input_ids"]
@@ -74,40 +65,38 @@ def classify_text(sentence: str):
 
     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")
+    
+    # Check if there are at least two words
+    word_count = len(user_input.split())
+    if word_count < 2:
+        raise HTTPException(status_code=400, detail="Text must contain at least two words")
 
-    # Run classification asynchronously to prevent blocking
-    result, perplexity = await asyncio.to_thread(classify_text, user_input)
 
+    result, perplexity = await asyncio.to_thread(classify_text, user_input)
+    
     return {
         "result": result,
         "perplexity": round(perplexity, 2),
     }
 
-
 # Health check route
-
-
 @app.get("/health")
 async def health_check():
     return {"status": "ok"}
 
-
 # Simple index route
-
-
 @app.get("/")
 def index():
     return {
         "message": "FastAPI API is up.",
         "try": "/docs to test the API.",
-        "status": "OK",
+        "status": "OK"
     }
+