| 480 | |
| 481 | |
| 482 | class KronosPredictor: |
| 483 | |
| 484 | def __init__(self, model, tokenizer, device=None, max_context=512, clip=5): |
| 485 | self.tokenizer = tokenizer |
| 486 | self.model = model |
| 487 | self.max_context = max_context |
| 488 | self.clip = clip |
| 489 | self.price_cols = ['open', 'high', 'low', 'close'] |
| 490 | self.vol_col = 'volume' |
| 491 | self.amt_vol = 'amount' |
| 492 | self.time_cols = ['minute', 'hour', 'weekday', 'day', 'month'] |
| 493 | |
| 494 | # Auto-detect device if not specified |
| 495 | if device is None: |
| 496 | if torch.cuda.is_available(): |
| 497 | device = "cuda:0" |
| 498 | elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): |
| 499 | device = "mps" |
| 500 | else: |
| 501 | device = "cpu" |
| 502 | |
| 503 | self.device = device |
| 504 | |
| 505 | self.tokenizer = self.tokenizer.to(self.device) |
| 506 | self.model = self.model.to(self.device) |
| 507 | |
| 508 | def generate(self, x, x_stamp, y_stamp, pred_len, T, top_k, top_p, sample_count, verbose): |
| 509 | |
| 510 | x_tensor = torch.from_numpy(np.array(x).astype(np.float32)).to(self.device) |
| 511 | x_stamp_tensor = torch.from_numpy(np.array(x_stamp).astype(np.float32)).to(self.device) |
| 512 | y_stamp_tensor = torch.from_numpy(np.array(y_stamp).astype(np.float32)).to(self.device) |
| 513 | |
| 514 | preds = auto_regressive_inference(self.tokenizer, self.model, x_tensor, x_stamp_tensor, y_stamp_tensor, self.max_context, pred_len, |
| 515 | self.clip, T, top_k, top_p, sample_count, verbose) |
| 516 | preds = preds[:, -pred_len:, :] |
| 517 | return preds |
| 518 | |
| 519 | def predict(self, df, x_timestamp, y_timestamp, pred_len, T=1.0, top_k=0, top_p=0.9, sample_count=1, verbose=True): |
| 520 | |
| 521 | if not isinstance(df, pd.DataFrame): |
| 522 | raise ValueError("Input must be a pandas DataFrame.") |
| 523 | |
| 524 | if not all(col in df.columns for col in self.price_cols): |
| 525 | raise ValueError(f"Price columns {self.price_cols} not found in DataFrame.") |
| 526 | |
| 527 | df = df.copy() |
| 528 | if self.vol_col not in df.columns: |
| 529 | df[self.vol_col] = 0.0 # Fill missing volume with zeros |
| 530 | df[self.amt_vol] = 0.0 # Fill missing amount with zeros |
| 531 | if self.amt_vol not in df.columns and self.vol_col in df.columns: |
| 532 | df[self.amt_vol] = df[self.vol_col] * df[self.price_cols].mean(axis=1) |
| 533 | |
| 534 | if df[self.price_cols + [self.vol_col, self.amt_vol]].isnull().values.any(): |
| 535 | raise ValueError("Input DataFrame contains NaN values in price or volume columns.") |
| 536 | |
| 537 | x_time_df = calc_time_stamps(x_timestamp) |
| 538 | y_time_df = calc_time_stamps(y_timestamp) |
| 539 |
no outgoing calls