ames_model / main.py
xValentim's picture
done
a233326
from fastapi import FastAPI
import pickle
import uvicorn
import pandas as pd
import json
import time
app = FastAPI(debug=True)
# Function to load pickle file
def load_pickle(filename):
with open(filename, 'rb') as file:
data = pickle.load(file)
return data
# Load pickle file
ml_model = load_pickle('my_model.pkl')
#Endpoints
#Root endpoints
@app.get("/")
def root():
return {"API": "An API for AMES Prediction."}
@app.post('/predict')
async def predict(data: dict):
start = time.time()
data_ = pd.DataFrame(data)
print('Data: ', type(data_))
model_output = ml_model.predict(data_)
print('Predictions: ', 10 ** model_output)
end = time.time()
time_to_inference = end - start
return {'predict': list(10 ** model_output), 'time_to_inference': time_to_inference}