FilesExpand file tree

master

/

prediction_example.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

80 lines (61 loc) · 2.42 KB

master

/

prediction_example.py

Copy path

Top

File metadata and controls

• Code
• Blame

80 lines (61 loc) · 2.42 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

import pandas as pd

import matplotlib.pyplot as plt

import sys

sys.path.append("../")

from model import Kronos, KronosTokenizer, KronosPredictor

def plot_prediction(kline_df, pred_df):

pred_df.index = kline_df.index[-pred_df.shape[0]:]

sr_close = kline_df['close']

sr_pred_close = pred_df['close']

sr_close.name = 'Ground Truth'

sr_pred_close.name = "Prediction"

sr_volume = kline_df['volume']

sr_pred_volume = pred_df['volume']

sr_volume.name = 'Ground Truth'

sr_pred_volume.name = "Prediction"

close_df = pd.concat([sr_close, sr_pred_close], axis=1)

volume_df = pd.concat([sr_volume, sr_pred_volume], axis=1)

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 6), sharex=True)

ax1.plot(close_df['Ground Truth'], label='Ground Truth', color='blue', linewidth=1.5)

ax1.plot(close_df['Prediction'], label='Prediction', color='red', linewidth=1.5)

ax1.set_ylabel('Close Price', fontsize=14)

ax1.legend(loc='lower left', fontsize=12)

ax1.grid(True)

ax2.plot(volume_df['Ground Truth'], label='Ground Truth', color='blue', linewidth=1.5)

ax2.plot(volume_df['Prediction'], label='Prediction', color='red', linewidth=1.5)

ax2.set_ylabel('Volume', fontsize=14)

ax2.legend(loc='upper left', fontsize=12)

ax2.grid(True)

plt.tight_layout()

plt.show()

# 1. Load Model and Tokenizer

tokenizer = KronosTokenizer.from_pretrained("NeoQuasar/Kronos-Tokenizer-base")

model = Kronos.from_pretrained("NeoQuasar/Kronos-small")

# 2. Instantiate Predictor

predictor = KronosPredictor(model, tokenizer, device="cuda:0", max_context=512)

# 3. Prepare Data

df = pd.read_csv("./data/XSHG_5min_600977.csv")

df['timestamps'] = pd.to_datetime(df['timestamps'])

lookback = 400

pred_len = 120

x_df = df.loc[:lookback-1, ['open', 'high', 'low', 'close', 'volume', 'amount']]

x_timestamp = df.loc[:lookback-1, 'timestamps']

y_timestamp = df.loc[lookback:lookback+pred_len-1, 'timestamps']

# 4. Make Prediction

pred_df = predictor.predict(

df=x_df,

x_timestamp=x_timestamp,

y_timestamp=y_timestamp,

pred_len=pred_len,

T=1.0,

top_p=0.9,

sample_count=1,

verbose=True

)

# 5. Visualize Results

print("Forecasted Data Head:")

print(pred_df.head())

# Combine historical and forecasted data for plotting

kline_df = df.loc[:lookback+pred_len-1]

# visualize

plot_prediction(kline_df, pred_df)