# Path: docker/payloads.py
# class CompletionPayload(BaseModel):
#     context: str
#     temp: Optional[float] = 1.0
#     top_p: Optional[float] = 0.9
#     top_k: Optional[int] = 50
#     length: Optional[int] = 256
# 
# class CompletionResponse(BaseModel):
#     context: str
#     completion: str
#     time: float
# 
# class QueueRequest(BaseModel):
#     qid: int
# 
# class QueueResponse(BaseModel):
#     qid: int
# 
# Path: docker/ops.py
# def get_gptj_model():
#     compile_model()
#     return gptj_model


import uvicorn
import traceback
import logging
from fastapi import FastAPI
from starlette.middleware.cors import CORSMiddleware
from .payloads import CompletionPayload, CompletionResponse, QueueRequest, QueueResponse
from .ops import get_gptj_model
    allow_methods=["*"],
    allow_headers=["*"],
)


# globals
MODEL_API = None


@app.on_event("startup")
async def startup_event():
    global MODEL_API
    try:
        MODEL_API = get_gptj_model()
        MODEL_API.load_model()
        MODEL_API.start_background()
    except Exception as e:
        logger.debug(f"Model could not be loaded: {str(e)}")
        traceback.print_exc()


@app.post("/model/get_prediction")
def get_prediction(payload: QueueRequest) -> CompletionResponse:
    res = MODEL_API.wait_for_queue(payload.qid)
    return CompletionResponse(**res)


@app.post("/model/predict")
def model_prediction(payload: CompletionPayload) -> QueueResponse:
    res = MODEL_API.add_to_queue(payload)

# Path: docker/payloads.py
# class CompletionPayload(BaseModel):
#     context: str
#     temp: Optional[float] = 1.0
#     top_p: Optional[float] = 0.9
#     top_k: Optional[int] = 50
#     length: Optional[int] = 256
# 
# class CompletionResponse(BaseModel):
#     context: str
#     completion: str
#     time: float
# 
# class QueueRequest(BaseModel):
#     qid: int
# 
# class QueueResponse(BaseModel):
#     qid: int
# 
# Path: docker/ops.py
# def get_gptj_model():
#     compile_model()
#     return gptj_model


# Path: docker/main.py
import uvicorn
import traceback
import logging
from fastapi import FastAPI
from starlette.middleware.cors import CORSMiddleware
from .payloads import CompletionPayload, CompletionResponse, QueueRequest, QueueResponse
from .ops import get_gptj_model

    allow_methods=["*"],
    allow_headers=["*"],
)


# globals
MODEL_API = None


@app.on_event("startup")
async def startup_event():
    global MODEL_API
    try:
        MODEL_API = get_gptj_model()
        MODEL_API.load_model()
        MODEL_API.start_background()
    except Exception as e:
        logger.debug(f"Model could not be loaded: {str(e)}")
        traceback.print_exc()


@app.post("/model/get_prediction")
def get_prediction(payload: QueueRequest) -> CompletionResponse:
    res = MODEL_API.wait_for_queue(payload.qid)
    return CompletionResponse(**res)


@app.post("/model/predict")
def model_prediction(payload: CompletionPayload) -> QueueResponse:
    res = MODEL_API.add_to_queue(payload)

    return QueueResponse(qid=res['qid'])