File size: 6,261 Bytes
91c60e6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 |
#!/usr/bin/env python3
"""
Secure FastAPI server for serving Mistral 7B fine-tuned models
Includes API key authentication like commercial services
"""
import os
import sys
import secrets
from typing import Optional
from fastapi import FastAPI, HTTPException, Header, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
import uvicorn
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from inference.inference_mistral7b import load_local_model, generate_with_local_model, get_device_info
import torch
# Configuration - Resolve model path relative to msp root
_MODEL_BASE = Path(__file__).parent.parent / "mistral7b-finetuned-ahb2apb"
DEFAULT_MODEL_PATH = str(_MODEL_BASE)
# API Key authentication
API_KEYS = set()
API_KEY_FILE = "api_keys.txt"
# Load or generate API keys
def load_api_keys():
"""Load API keys from file or create default"""
global API_KEYS
if os.path.exists(API_KEY_FILE):
with open(API_KEY_FILE, 'r') as f:
API_KEYS = {line.strip() for line in f if line.strip()}
else:
# Generate a default API key
default_key = secrets.token_urlsafe(32)
with open(API_KEY_FILE, 'w') as f:
f.write(default_key + '\n')
API_KEYS = {default_key}
print(f"\n🔑 Generated default API key: {default_key}")
print(f" Save this key! Store it in: {API_KEY_FILE}")
print(f"✓ Loaded {len(API_KEYS)} API key(s)")
def verify_api_key(api_key: str = Header(None)):
"""Verify API key in request header"""
if api_key is None:
raise HTTPException(
status_code=401,
detail="API key required. Add header: 'X-API-Key: your-api-key'"
)
if api_key not in API_KEYS:
raise HTTPException(status_code=403, detail="Invalid API key")
return api_key
# Global model and tokenizer
model = None
tokenizer = None
device_info = None
app = FastAPI(
title="Mistral 7B AHB2APB API (Secure)",
description="Secure API for serving the fine-tuned Mistral 7B model for AHB2APB conversion",
version="1.0.0"
)
# Enable CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # In production, restrict this!
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Security scheme
security = HTTPBearer(auto_error=False)
# Request/Response models
class GenerateRequest(BaseModel):
prompt: str
max_length: Optional[int] = 512
temperature: Optional[float] = 0.7
class GenerateResponse(BaseModel):
response: str
model: str
max_length: int
temperature: float
class HealthResponse(BaseModel):
status: str
model_loaded: bool
device: str
model_path: str
authentication: str
@app.on_event("startup")
async def startup():
"""Load model and API keys on startup"""
global model, tokenizer, device_info
load_api_keys()
model_path = os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH)
print(f"\nLoading model from: {model_path}")
print("=" * 70)
try:
device_info = get_device_info()
model, tokenizer = load_local_model(model_path)
print(f"\n✓ Model loaded successfully on {device_info['device']}!")
print(f"✓ API server ready (authentication enabled)")
print("=" * 70)
except Exception as e:
print(f"\n✗ Error loading model: {e}")
sys.exit(1)
@app.get("/health", response_model=HealthResponse)
async def health_check():
"""Health check endpoint (no auth required)"""
return HealthResponse(
status="healthy" if model is not None else "error",
model_loaded=model is not None,
device=device_info["device"] if device_info else "unknown",
model_path=os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH),
authentication="enabled"
)
@app.get("/")
async def root():
"""Root endpoint with API information"""
return {
"name": "Mistral 7B AHB2APB API (Secure)",
"version": "1.0.0",
"status": "running",
"authentication": "API key required",
"model": os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH),
"endpoints": {
"health": "/health",
"generate": "/api/generate (requires API key)",
"docs": "/docs"
}
}
@app.post("/api/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest, api_key: str = Depends(verify_api_key)):
"""Generate text from a prompt (requires API key)"""
if model is None or tokenizer is None:
raise HTTPException(status_code=503, detail="Model not loaded")
try:
response = generate_with_local_model(
model=model,
tokenizer=tokenizer,
prompt=request.prompt,
max_length=request.max_length or 512,
temperature=request.temperature or 0.7
)
return GenerateResponse(
response=response,
model=os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH),
max_length=request.max_length or 512,
temperature=request.temperature or 0.7
)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Generation error: {str(e)}")
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Start Secure Mistral 7B API server")
parser.add_argument("--model-path", type=str, default=DEFAULT_MODEL_PATH)
parser.add_argument("--host", type=str, default="0.0.0.0")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--reload", action="store_true")
args = parser.parse_args()
os.environ["MODEL_PATH"] = args.model_path
print(f"\n🔒 Starting Secure Mistral 7B AHB2APB API Server")
print(f" Model: {args.model_path}")
print(f" Host: {args.host}")
print(f" Port: {args.port}\n")
# Change to api directory for proper module resolution
import os
os.chdir(os.path.dirname(os.path.abspath(__file__)))
uvicorn.run("api_server_secure:app", host=args.host, port=args.port, reload=args.reload)
|