Merge pull request #10 from cyberalertnepal/PujanDev

feat: reduce load while model is first processing the data

app.py

@@ -6,11 +6,7 @@ import nltk
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Ensure punkt is available
-    nltk.download("punkt")
 
-    nltk.download('punkt_tab')
-    
     # Your model warmup
     warmup()
     yield

features/text_classifier/model_loader.py

@@ -5,7 +5,7 @@ from transformers import GPT2LMHeadModel, GPT2TokenizerFast, GPT2Config
 from huggingface_hub import snapshot_download
 import torch
 from dotenv import load_dotenv
-
+import nltk
 load_dotenv()
 REPO_ID = "Pujan-Dev/AI-Text-Detector"
 MODEL_DIR = "./models"
@@ -18,6 +18,11 @@ _model, _tokenizer = None, None
 
 def warmup():
     global _model, _tokenizer
+    # Ensure punkt is available
+    nltk.download("punkt")
+
+    nltk.download('punkt_tab')
+
     download_model_repo()
     _model, _tokenizer = load_model()
     logging.info("Its ready")

features/text_classifier/preprocess.py

@@ -3,8 +3,6 @@ import docx
 from io import BytesIO
 import logging
 from fastapi import HTTPException
-
-
 def parse_docx(file: BytesIO):
     doc = docx.Document(file)
     text = ""
@@ -23,7 +21,8 @@ def parse_pdf(file: BytesIO):
         return text
     except Exception as e:
         logging.error(f"Error while processing PDF: {str(e)}")
-        raise HTTPException(status_code=500, detail="Error processing PDF file")
+        raise HTTPException(
+            status_code=500, detail="Error processing PDF file")
 
 
 def parse_txt(file: BytesIO):

features/text_classifier/routes.py

@@ -1,38 +1,43 @@
-from fastapi import APIRouter, Depends, HTTPException, UploadFile, File
-from fastapi.security import HTTPBearer
+from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, status
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
-from .controller import handle_text_analysis, handle_file_upload, handle_sentence_level_analysis,handle_file_sentance
+from .controller import handle_text_analysis, handle_file_upload, handle_sentence_level_analysis, handle_file_sentance
+import os
 
 router = APIRouter()
-bearer_scheme = HTTPBearer()
-
+security = HTTPBearer()
 
 class TextInput(BaseModel):
     text: str
 
+async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    token = credentials.credentials
+    if token != os.getenv("MY_SECRET_TOKEN"):  # Replace with your actual secret
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Invalid or expired token"
+        )
+    return token
 
 @router.post("/analyse")
-async def analyze(data: TextInput, token: str = Depends(bearer_scheme)):
+async def analyze(data: TextInput, token: str = Depends(verify_token)):
     return await handle_text_analysis(data.text)
 
-
 @router.post("/upload")
-async def upload_file(
-    file: UploadFile = File(...), token: str = Depends(bearer_scheme)
-):
+async def upload_file(file: UploadFile = File(...), token: str = Depends(verify_token)):
     return await handle_file_upload(file)
 
-
-@router.get("/health")
-def health():
-    return {"status": "ok"}
-
 @router.post("/analyse-sentences")
-async def analyze_sentences(data: TextInput, token: str = Depends(bearer_scheme)):
+async def analyze_sentences(data: TextInput, token: str = Depends(verify_token)):
     if not data.text:
         raise HTTPException(status_code=400, detail="Missing 'text' in request body")
     return await handle_sentence_level_analysis(data.text)
 
 @router.post("/analyse-sentance-file")
-async def AnalyzeSentanceFile(file:UploadFile=File(...),token:str=Depends(bearer_scheme)):
+async def AnalyzeSentanceFile(file: UploadFile = File(...), token: str = Depends(verify_token)):
     return await handle_file_sentance(file)
+
+@router.get("/health")
+def health():
+    return {"status": "ok"}
+