Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 import torch
 import secrets
 import time
+import json
 from fastapi import FastAPI, HTTPException, Security, Depends
 from fastapi.security.api_key import APIKeyHeader
 from pydantic import BaseModel
@@ -13,62 +14,98 @@ 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. 
-# Note: These will reset if the Space restarts unless you use Persistent Storage.
+# In-memory storage for keys
+# Note: On a free Space, these reset on restart. 
 generated_keys = {} 
 
-app = FastAPI(title="Overflow-111.7B API Manager")
+app = FastAPI(title="Overflow-111.7B API")
 
 # --- MODEL LOADING ---
-print("Loading Overflow-111.7B Engine...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_PATH,
-    trust_remote_code=True,
-    device_map={"": "cpu"},
-    torch_dtype=torch.bfloat16,
-    low_cpu_mem_usage=True
-)
+# We load these at the global level so they stay in memory
+print("Starting Engine: Loading Overflow-111.7B (1-Bit Logic)...")
 
+try:
+    # 1. Load Tokenizer - Explicitly trust remote code to avoid build hangs
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_PATH, 
+        trust_remote_code=True
+    )
+
+    # 2. Load Model - Optimized for 16GB CPU RAM
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_PATH,
+        trust_remote_code=True,
+        device_map={"": "cpu"},
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True
+    )
+    print("Engine Status: ONLINE")
+except Exception as e:
+    print(f"CRITICAL LOADING ERROR: {e}")
+
+# --- SCHEMAS ---
 class Query(BaseModel):
     prompt: str
     max_tokens: int = 50
+    temperature: float = 0.7
 
-# --- API KEY GENERATION ---
+# --- API KEY LOGIC ---
 @app.get("/api/generate")
 async def create_new_key():
-    """Generates a new of_sk key for the user."""
-    # Generate a random string of 24 characters
-    random_hex = secrets.token_hex(12)
-    new_key = f"of_sk-{random_hex}"
-    
-    # Store with a timestamp
+    """Generates a unique of_sk- key for the user."""
+    new_key = f"of_sk-{secrets.token_hex(12)}"
     generated_keys[new_key] = {"created_at": time.time()}
-    
     return {
         "status": "success",
         "api_key": new_key,
-        "instructions": f"Include this key in your request header as '{API_KEY_NAME}'"
+        "instructions": f"Add this to your headers as '{API_KEY_NAME}'"
     }
 
-# --- SECURITY CHECK ---
 async def get_api_key(api_key_header: str = Depends(api_key_header)):
     if api_key_header in generated_keys:
         return api_key_header
+    # Also allow a master key from environment variables if set
+    master_key = os.environ.get("MASTER_API_KEY")
+    if master_key and api_key_header == master_key:
+        return api_key_header
+        
     raise HTTPException(
         status_code=HTTP_403_FORBIDDEN, 
-        detail="Invalid or expired API Key. Generate one at /api/generate"
+        detail="Invalid API Key. Generate one at /api/generate"
     )
 
+# --- ENDPOINTS ---
 @app.post("/v1/generate")
 async def generate(query: Query, api_key: str = Depends(get_api_key)):
-    inputs = tokenizer(query.prompt, return_tensors="pt")
-    with torch.no_grad():
-        output_tokens = model.generate(**inputs, max_new_tokens=query.max_tokens)
-    
-    response = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
-    return {"text": response}
+    try:
+        inputs = tokenizer(query.prompt, return_tensors="pt")
+        
+        with torch.no_grad():
+            output_tokens = model.generate(
+                **inputs, 
+                max_new_tokens=query.max_tokens,
+                temperature=query.temperature,
+                do_sample=True if query.temperature > 0 else False
+            )
+        
+        response_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+        return {
+            "object": "text_completion",
+            "model": "Overflow-111.7B",
+            "choices": [{"text": response_text}]
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")
-def home():
-    return {"message": "Welcome to Overflow-111.7B. Go to /api/generate to get a key."}
+def health_check():
+    return {
+        "status": "active",
+        "model": "Overflow-111.7B",
+        "device": "CPU",
+        "usage": "Go to /api/generate to get started"
+    }
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)