Update app.py

app.py

@@ -1,8 +1,9 @@
 import os
+import sys
 import torch
 import secrets
 import time
-from fastapi import FastAPI, HTTPException, Security, Depends
+from fastapi import FastAPI, HTTPException, Depends
 from fastapi.security.api_key import APIKeyHeader
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
@@ -13,7 +14,7 @@ MODEL_PATH = "/app/model"
 API_KEY_NAME = "X-API-Key"
 api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
 
-# In-memory storage for keys generated via /api/generate
+# In-memory storage for keys
 generated_keys = {} 
 
 app = FastAPI(title="Overflow-111.7B API")
@@ -22,19 +23,22 @@ app = FastAPI(title="Overflow-111.7B API")
 print("Starting Engine: Loading Overflow-111.7B (1-Bit Logic)...")
 
 try:
-    # 1. Register and Load Custom Config
-    # This fixes the "Unrecognized configuration class" error
+    # 1. Inject model path into system path so Python can find custom modules
+    if MODEL_PATH not in sys.path:
+        sys.path.append(MODEL_PATH)
+
+    # 2. Register the custom config class
+    # This assumes the file in /app/model is named 'configuration_overflow.py'
     from configuration_overflow import OverflowConfig
     AutoConfig.register("overflow", OverflowConfig)
 
-    # 2. Load Tokenizer
-    # We use trust_remote_code and the registered config to load smoothly
+    # 3. Load Tokenizer
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_PATH, 
         trust_remote_code=True
     )
 
-    # 3. Load Model - Optimized for 16GB CPU RAM
+    # 4. Load Model
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_PATH,
         trust_remote_code=True,
@@ -52,7 +56,7 @@ class Query(BaseModel):
     max_tokens: int = 50
     temperature: float = 0.7
 
-# --- API KEY LOGIC ---
+# --- AUTH LOGIC ---
 @app.get("/api/generate")
 async def create_new_key():
     """Generates a unique of_sk- key."""
@@ -64,16 +68,14 @@ async def create_new_key():
         "instructions": f"Add this to your headers as '{API_KEY_NAME}'"
     }
 
-async def get_api_key(api_key_header: str = Depends(api_key_header)):
-    # Check in-memory keys
-    if api_key_header in generated_keys:
-        return api_key_header
+async def verify_auth(api_key: str = Depends(api_key_header)):
+    if api_key in generated_keys:
+        return api_key
     
-    # Check for a MASTER_KEY set in Space Secrets (Settings > Variables and Secrets)
-    # Use this if you don't want to keep generating new keys
+    # Check for a MASTER_API_KEY set in Space Secrets/Variables
     master_key = os.environ.get("MASTER_API_KEY")
-    if master_key and api_key_header == master_key:
-        return api_key_header
+    if master_key and api_key == master_key:
+        return api_key
         
     raise HTTPException(
         status_code=HTTP_403_FORBIDDEN, 
@@ -82,7 +84,7 @@ async def get_api_key(api_key_header: str = Depends(api_key_header)):
 
 # --- ENDPOINTS ---
 @app.post("/v1/generate")
-async def generate(query: Query, api_key: str = Depends(get_api_key)):
+async def generate(query: Query, auth: str = Depends(verify_auth)):
     try:
         inputs = tokenizer(query.prompt, return_tensors="pt")
         
@@ -108,8 +110,7 @@ def health_check():
     return {
         "status": "active",
         "model": "Overflow-111.7B",
-        "device": "CPU",
-        "info": "Visit /api/generate to get an API key."
+        "auth_method": "X-API-Key"
     }
 
 if __name__ == "__main__":