Upload deployment/api_server.py with huggingface_hub

deployment/api_server.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python3
+"""
+AuraMind REST API Server
+Production-ready API for AuraMind smartphone deployment
+"""
+
+from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional, List, Dict
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import uvicorn
+import logging
+import time
+from datetime import datetime
+import os
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Request/Response models
+class ChatRequest(BaseModel):
+    message: str
+    mode: str = "Assistant"  # "Therapist" or "Assistant"
+    max_tokens: int = 200
+    temperature: float = 0.7
+
+class ChatResponse(BaseModel):
+    response: str
+    mode: str
+    inference_time_ms: float
+    timestamp: str
+
+class ModelInfo(BaseModel):
+    variant: str
+    memory_usage: str
+    inference_speed: str
+    status: str
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="AuraMind API",
+    description="Smartphone-optimized dual-mode AI companion API",
+    version="1.0.0"
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Configure appropriately for production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Global model variables
+tokenizer = None
+model = None
+model_variant = None
+
+def load_model(variant: str = "270m"):
+    """Load AuraMind model"""
+    global tokenizer, model, model_variant
+    
+    try:
+        logger.info(f"Loading AuraMind {variant}...")
+        
+        model_name = "zail-ai/Auramind"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            low_cpu_mem_usage=True
+        )
+        
+        model.eval()
+        model_variant = variant
+        
+        logger.info(f"✅ AuraMind {variant} loaded successfully")
+        
+    except Exception as e:
+        logger.error(f"Failed to load model: {e}")
+        raise
+
+@app.on_event("startup")
+async def startup_event():
+    """Initialize model on startup"""
+    variant = os.getenv("MODEL_VARIANT", "270m")
+    load_model(variant)
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "model_loaded": model is not None,
+        "variant": model_variant,
+        "timestamp": datetime.now().isoformat()
+    }
+
+@app.get("/model/info", response_model=ModelInfo)
+async def get_model_info():
+    """Get model information"""
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    variant_configs = {
+        "270m": {"memory": "~680MB RAM", "speed": "100-300ms"},
+        "180m": {"memory": "~450MB RAM", "speed": "80-200ms"},
+        "90m": {"memory": "~225MB RAM", "speed": "50-150ms"}
+    }
+    
+    config = variant_configs.get(model_variant, {"memory": "Unknown", "speed": "Unknown"})
+    
+    return ModelInfo(
+        variant=model_variant,
+        memory_usage=config["memory"],
+        inference_speed=config["speed"],
+        status="ready"
+    )
+
+@app.post("/chat", response_model=ChatResponse)
+async def chat(request: ChatRequest):
+    """Generate chat response"""
+    if model is None or tokenizer is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    if request.mode not in ["Therapist", "Assistant"]:
+        raise HTTPException(status_code=400, detail="Mode must be 'Therapist' or 'Assistant'")
+    
+    try:
+        start_time = time.time()
+        
+        # Format prompt
+        prompt = f"<|start_of_turn|>user\n[{request.mode} Mode] {request.message}<|end_of_turn|>\n<|start_of_turn|>model\n"
+        
+        # Tokenize
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512
+        )
+        
+        # Generate
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=request.max_tokens,
+                temperature=request.temperature,
+                do_sample=True,
+                top_p=0.9,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode response
+        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = full_response.split("<|start_of_turn|>model\n")[-1].strip()
+        
+        inference_time = (time.time() - start_time) * 1000
+        
+        return ChatResponse(
+            response=response,
+            mode=request.mode,
+            inference_time_ms=round(inference_time, 2),
+            timestamp=datetime.now().isoformat()
+        )
+        
+    except Exception as e:
+        logger.error(f"Error generating response: {e}")
+        raise HTTPException(status_code=500, detail="Failed to generate response")
+
+@app.post("/chat/batch")
+async def chat_batch(requests: List[ChatRequest]):
+    """Process multiple chat requests"""
+    if len(requests) > 10:  # Limit batch size
+        raise HTTPException(status_code=400, detail="Batch size limited to 10 requests")
+    
+    responses = []
+    for req in requests:
+        response = await chat(req)
+        responses.append(response)
+    
+    return {"responses": responses}
+
+if __name__ == "__main__":
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=int(os.getenv("PORT", 8000)),
+        workers=1  # Single worker for model consistency
+    )