Update app.py

app.py

@@ -8,13 +8,16 @@ from fastapi import FastAPI, HTTPException, Depends
 from fastapi.security.api_key import APIKeyHeader
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
-from starlette.status import HTTP_403_FORBIDDEN
+from starlette.status import HTTP_403_FORBIDDEN, HTTP_503_SERVICE_UNAVAILABLE
 
 # --- CONFIGURATION ---
 MODEL_PATH = "/app/model"
 API_KEY_NAME = "X-API-Key"
 api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
 
+# Initialize variables at the global level so the functions can see them
+tokenizer = None
+model = None
 generated_keys = {} 
 
 app = FastAPI(title="Overflow-111.7B Self-Registering API")
@@ -28,31 +31,29 @@ try:
         sys.path.insert(0, MODEL_PATH)
 
     # 2. FORCE-REGISTER CONFIGURATION
-    # We manually import the file to ensure the class is loaded into memory
     import configuration_overflow
-    # Use the exact class name you provided
     conf_class = configuration_overflow.OverflowConfig
     AutoConfig.register("overflow", conf_class)
     print(f"Successfully registered 'overflow' config.")
 
     # 3. FORCE-REGISTER MODEL
-    # We need to find the model class in modeling_overflow.py
     import modeling_overflow
-    # Search for the class that ends with 'ForCausalLM'
     model_classes = [c for c in dir(modeling_overflow) if 'ForCausalLM' in c]
     if model_classes:
         model_class = getattr(modeling_overflow, model_classes[0])
         AutoModelForCausalLM.register(conf_class, model_class)
         print(f"Successfully registered {model_classes[0]} to AutoModel.")
 
-    # 4. LOAD TOKENIZER
+    # 4. LOAD TOKENIZER & MODEL
+    # We use 'global' to update the variables we defined at the top
+    global tokenizer, model
+    
     print("Loading Tokenizer...")
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_PATH, 
         trust_remote_code=True
     )
 
-    # 5. LOAD MODEL (CPU Optimized)
     print("Loading Model Weights (111.7B Parameters - 1-Bit)...")
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_PATH,
@@ -87,6 +88,13 @@ async def verify_auth(api_key: str = Depends(api_key_header)):
 # --- ENDPOINTS ---
 @app.post("/v1/generate")
 async def generate(query: Query, auth: str = Depends(verify_auth)):
+    # Safety Check: If the model hasn't finished loading yet
+    if tokenizer is None or model is None:
+        raise HTTPException(
+            status_code=HTTP_503_SERVICE_UNAVAILABLE, 
+            detail="Engine is still booting up. Please wait a moment."
+        )
+
     try:
         inputs = tokenizer(query.prompt, return_tensors="pt")
         with torch.no_grad():
@@ -103,7 +111,8 @@ async def generate(query: Query, auth: str = Depends(verify_auth)):
 
 @app.get("/")
 def health():
-    return {"status": "active", "engine": "Overflow-111.7B"}
+    state = "active" if model else "loading"
+    return {"status": state, "engine": "Overflow-111.7B"}
 
 if __name__ == "__main__":
     import uvicorn