Upload models/msp/api/api_server.py with huggingface_hub

models/msp/api/api_server.py

@@ -0,0 +1,216 @@
+#!/usr/bin/env python3
+"""
+FastAPI server for serving Mistral 7B fine-tuned models
+"""
+
+import os
+import sys
+from typing import Optional, Dict, Any
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import uvicorn
+import sys
+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)
+
+# Global model and tokenizer (loaded once at startup)
+model = None
+tokenizer = None
+device_info = None
+
+app = FastAPI(
+    title="Mistral 7B AHB2APB API",
+    description="API for serving the fine-tuned Mistral 7B model for AHB2APB conversion",
+    version="1.0.0"
+)
+
+# Enable CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# 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
+
+@app.on_event("startup")
+async def load_model():
+    """Load the model when the server starts"""
+    global model, tokenizer, device_info
+    
+    model_path = os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH)
+    
+    print(f"Loading 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"✓ Server ready to accept requests")
+        print("=" * 70)
+    except Exception as e:
+        print(f"\n✗ Error loading model: {e}")
+        print("=" * 70)
+        sys.exit(1)
+
+@app.get("/health", response_model=HealthResponse)
+async def health_check():
+    """Health check endpoint"""
+    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)
+    )
+
+@app.get("/")
+async def root():
+    """Root endpoint with API information"""
+    return {
+        "name": "Mistral 7B AHB2APB API",
+        "version": "1.0.0",
+        "status": "running",
+        "model": os.environ.get("MODEL_PATH", DEFAULT_MODEL_PATH),
+        "endpoints": {
+            "health": "/health",
+            "generate": "/api/generate",
+            "docs": "/docs"
+        }
+    }
+
+@app.post("/api/generate", response_model=GenerateResponse)
+async def generate(request: GenerateRequest):
+    """
+    Generate text from a prompt using the fine-tuned model
+    """
+    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)}")
+
+@app.post("/api/generate/batch")
+async def generate_batch(requests: list[GenerateRequest]):
+    """
+    Generate text from multiple prompts (batch processing)
+    """
+    if model is None or tokenizer is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        responses = []
+        for req in requests:
+            response = generate_with_local_model(
+                model=model,
+                tokenizer=tokenizer,
+                prompt=req.prompt,
+                max_length=req.max_length or 512,
+                temperature=req.temperature or 0.7
+            )
+            responses.append({
+                "response": response,
+                "prompt": req.prompt
+            })
+        
+        return {"results": responses}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Batch generation error: {str(e)}")
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Start Mistral 7B API server")
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        default=DEFAULT_MODEL_PATH,
+        help=f"Path to fine-tuned model (default: {DEFAULT_MODEL_PATH})"
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="0.0.0.0",
+        help="Host to bind to (default: 0.0.0.0)"
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="Port to bind to (default: 8000)"
+    )
+    parser.add_argument(
+        "--reload",
+        action="store_true",
+        help="Enable auto-reload (for development)"
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help="Number of worker processes (default: 1)"
+    )
+    
+    args = parser.parse_args()
+    
+    # Set model path as environment variable for the startup event
+    os.environ["MODEL_PATH"] = args.model_path
+    
+    print(f"\n🚀 Starting Mistral 7B AHB2APB API Server")
+    print(f"   Model: {args.model_path}")
+    print(f"   Host: {args.host}")
+    print(f"   Port: {args.port}")
+    print(f"   Workers: {args.workers}")
+    print(f"   Reload: {args.reload}\n")
+    
+    # Change to api directory for proper module resolution
+    import os
+    os.chdir(os.path.dirname(os.path.abspath(__file__)))
+    uvicorn.run(
+        "api_server:app",
+        host=args.host,
+        port=args.port,
+        reload=args.reload,
+        workers=1 if args.reload else args.workers
+    )
+