File size: 2,016 Bytes
b81150d | 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 | from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import mlx_lm
import os
import time
from typing import List, Optional
from pydantic import BaseModel
# --- Fast, Apple Silicon Native Inference ---
# Uses mlx-lm library which is highly optimized for 16GB Mac RAM.
app = FastAPI(title="Ankahi Mac Server", version="1.0.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
MODEL_PATH = "ankahi_mlx_4bit" # The 4-bit optimized model folder
class PredictRequest(BaseModel):
pictogram_sequence: List[str]
persona_id: str
n_alternatives: int = 3
model = None
tokenizer = None
@app.on_event("startup")
async def load_optimized_model():
global model, tokenizer
print(f"Loading 4-bit optimized model from {MODEL_PATH}...")
try:
model, tokenizer = mlx_lm.load(MODEL_PATH)
print("Optimized model loaded on Apple Silicon!")
except Exception as e:
print(f"Error: Could not load model. Did you run convert_to_mlx.py first? {e}")
@app.post("/predict")
async def predict(request: PredictRequest):
if not model:
raise HTTPException(status_code=503, detail="Model not loaded")
start_time = time.time()
prompt = f"Pictogram sequence: [{', '.join(request.pictogram_sequence)}]. Predict {request.n_alternatives} sentences."
# Fast Apple Silicon Inference
response = mlx_lm.generate(model, tokenizer, prompt=prompt, max_tokens=150)
# Simple list parsing for demo (real version uses regex)
alternatives = [line.strip() for line in response.split('\n') if line.strip()][:request.n_alternatives]
return {
"alternatives": alternatives or [response],
"primary": alternatives[0] if alternatives else response,
"inference_time_ms": (time.time() - start_time) * 1000,
"persona_id": request.persona_id
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
|