File size: 550 Bytes
9433701
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from typing import Any, Dict
import os

from src.elliptic_fraud.predictor import FraudPredictor

class EndpointHandler:
    def __init__(self, path: str = "."):
        default_in = int(os.environ.get("IN_FEATURES", "4"))
        model_type = os.environ.get("MODEL_TYPE", "mlp")
        ckpt = os.environ.get("CKPT_PATH", "")
        self.predictor = FraudPredictor(default_in_features=default_in, model_type=model_type, ckpt_path=ckpt)

    def __call__(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        return self.predictor.predict(inputs)