Fix handler: use native Kronos model instead of ChronosPipeline

handler.py

@@ -1,26 +1,53 @@
 import torch
-from chronos import ChronosPipeline
+import json
+from safetensors.torch import load_file
+from model.kronos import Kronos
 
 class EndpointHandler:
     def __init__(self, path=""):
-        self.pipeline = ChronosPipeline.from_pretrained(
-            path,
-            device_map="auto",
-            torch_dtype=torch.float32,
+        with open(f"{path}/config.json", "r") as f:
+            config = json.load(f)
+
+        self.model = Kronos(
+            input_dim=config.get("input_dim", 5),
+            d_model=config.get("d_model", 256),
+            nhead=config.get("nhead", 8),
+            num_layers=config.get("num_layers", 6),
+            dim_feedforward=config.get("dim_feedforward", 1024),
+            max_seq_len=config.get("max_seq_len", 512),
+            output_dim=config.get("output_dim", 5),
+            dropout=config.get("dropout", 0.1),
         )
 
+        weights = load_file(f"{path}/model.safetensors")
+        self.model.load_state_dict(weights)
+        self.model.eval()
+
     def __call__(self, data):
         inputs = data.get("inputs", [])
         parameters = data.get("parameters", {})
         prediction_length = parameters.get("prediction_length", 8)
 
         if isinstance(inputs[0], list):
-            closes = [candle[3] for candle in inputs]
+            ohlcv = inputs
         else:
-            closes = inputs
+            ohlcv = [[v, v, v, v, 0] for v in inputs]
+
+        tensor = torch.tensor([ohlcv], dtype=torch.float32)
+
+        last_close = ohlcv[-1][3]
+        if last_close > 0:
+            tensor = tensor / last_close
+
+        with torch.no_grad():
+            output = self.model(tensor)
+
+        predicted = output[0, -prediction_length:, :].tolist()
 
-        context = torch.tensor([closes], dtype=torch.float32)
-        forecast = self.pipeline.predict(context, prediction_length=prediction_length, num_samples=20)
-        median_forecast = forecast.median(dim=1).values[0].tolist()
+        if last_close > 0:
+            predicted = [[v * last_close for v in candle] for candle in predicted]
 
-        return {"predictions": median_forecast}
+        return {
+            "predictions": predicted,
+            "prediction_length": prediction_length,
+        }

requirements.txt

@@ -1,2 +1,3 @@
-torch
-chronos-forecasting
+torch>=2.0.0
+safetensors
+einops