Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, HTTPException, status
+from fastapi import FastAPI, HTTPException, status, APIRouter, Request
 from pydantic import BaseModel, ValidationError
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 import torch
@@ -14,6 +14,10 @@ app = FastAPI(
     version="1.0.0"
 )
 
+# Create an API Router to manage endpoints.
+# This approach can sometimes help with routing issues in proxied environments.
+api_router = APIRouter()
+
 # Load model globally to avoid reloading on each request
 # This block runs once when the FastAPI application starts.
 try:
@@ -24,29 +28,26 @@ try:
     logger.info("Model loaded successfully.")
 except Exception as e:
     logger.exception("Failed to load model or tokenizer during startup!")
-    # Depending on the deployment, you might want to raise an exception here
-    # to prevent the app from starting if the model can't be loaded.
-    # For now, we'll let it potentially start and fail on prediction.
     raise RuntimeError(f"Could not load model: {e}")
 
 class InferenceRequest(BaseModel):
     """
-    Request model for the /predict endpoint.
+    Request model for the main prediction endpoint.
     Expects a single string field 'text' containing the sentence with [MASK] tokens.
     """
     text: str
 
 class PredictionResult(BaseModel):
     """
-    Response model for individual predictions from the /predict endpoint.
+    Response model for individual predictions from the API.
     """
     sequence: str    # The full sequence with the predicted token filled in
     score: float     # Confidence score of the prediction
     token: int       # The ID of the predicted token
     token_str: str   # The string representation of the predicted token
 
-@app.post(
-    "/predict",
+@api_router.post(
+    "/",  # Changed from "/predict" to "/" to funnel all POST requests to the root
     response_model=list[PredictionResult],
     summary="Predicts masked tokens in a given text",
     description="Accepts a text string with '[MASK]' tokens and returns top 5 predictions for each masked position."
@@ -74,7 +75,7 @@ async def predict_masked_lm(request: InferenceRequest):
         masked_token_indices = torch.where(inputs["input_ids"] == masked_token_id)[1]
 
         if not masked_token_indices.numel():
-            logger.warning("No [MASK] token found in the input text.")
+            logger.warning("No [MASK] token found in the input text. Returning 400 Bad Request.")
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail="Input text must contain at least one '[MASK]' token."
@@ -129,22 +130,32 @@ async def predict_masked_lm(request: InferenceRequest):
             detail=f"An internal server error occurred: {e}"
         )
 
-@app.get(
-    "/",
+@api_router.get(
+    "/health", # Health check moved to /health
     summary="Health Check",
     description="Returns a simple message indicating the API is running."
 )
-async def root():
+async def health_check():
     """
     Provides a basic health check endpoint for the API.
     """
     logger.info("Health check endpoint accessed.")
     return {"message": "NeuroBERT-Tiny API is running!"}
 
+# Include the API router in the main FastAPI application
+app.include_router(api_router)
+
+# Optional: Add a catch-all route for any unhandled paths.
+# This can help log when requests are hitting the app but to an unknown path.
+# This should now catch anything not / or /health
+@app.api_route("/{path_name:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
+async def catch_all(request: Request, path_name: str):
+    logger.warning(f"Unhandled route accessed: {request.method} {path_name}")
+    raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Not Found")
+
+
 # This block is for running the app directly, typically used for local development.
 # In a Docker container, Uvicorn (or Gunicorn) is usually invoked via the CMD in Dockerfile.
 if __name__ == "__main__":
     import uvicorn
-    # The 'reload=True' is great for local development for auto-reloading changes.
-    # For production in a Docker container, it's typically omitted for performance.
-    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
+    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")