# Ref. [Deploying a Machine Learning Model to AWS Lambda | TestDriven.io](https://testdriven.io/blog/ml-model-aws-lambda/)
import logging
from app import DocumentParserModel

# initialize logger and model during Lambda's cold start
LOGGER = logging.getLogger()
LOGGER.setLevel(logging.INFO)

model_path = "captcha.onnx"
img_size = (32, 128)
charset = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
model = DocumentParserModel(
    model_path=model_path,
    img_size=img_size,
    charset=charset,
)


def lambda_handle(event, context):
    # Only used to keep the Lambda warm
    if event.get("source") == "KEEP_LAMBDA_WARM":
        LOGGER.info("No ML work to do. Just staying warm...")
        return "Keeping Lambda warm"

    return {"statusCode": 200, "vc": model.predict_text(image_path=event["image_path"])}