Cerebrium script

ceremrium_local.py

# This script allows to simulate Cerebrium serverless in local for test purpose
# Author: @chaignc
# run me with
# uvicorn local:app --port 9000
  
from fastapi import FastAPI
from main import predict, PredictInput
import logging
import time

app = FastAPI()
logger = logging.basicConfig(level=logging.INFO)

# Define the predict endpoint
@app.post("/predict")
def predict_route(input_data: PredictInput):
    start_time = time.time()
    result = predict(input_data.dict(), "local", logger)
    end_time = time.time()
    run_time_ms = round((end_time - start_time) * 1000, 2)
    return {"result": result, "run_time_ms": run_time_ms}