| 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 |
|
|
| |
| |
|
|
| 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" |
|
|
| 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." |
| |
| |
| response = mlx_lm.generate(model, tokenizer, prompt=prompt, max_tokens=150) |
| |
| |
| 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) |
|
|
