File size: 3,749 Bytes
b8086d5
a90bc0e
6bf1eb6
 
 
 
 
 
 
 
 
 
b8086d5
 
 
a90bc0e
 
 
b8086d5
 
 
29091b3
6bf1eb6
 
 
 
b8086d5
a90bc0e
b8086d5
88fcfd5
a90bc0e
 
b8086d5
a90bc0e
b8086d5
 
a90bc0e
b8086d5
a90bc0e
b8086d5
 
1a10128
6bf1eb6
 
 
 
b8086d5
9927daa
6bf1eb6
b8086d5
a90bc0e
 
 
 
 
 
679c607
9b29694
 
 
 
 
 
a90bc0e
 
b8086d5
 
6bf1eb6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import numpy as np
import torch
import warnings

# Make chronos import optional
try:
    from chronos import BaseChronosPipeline
    CHRONOS_AVAILABLE = True
except ImportError:
    warnings.warn("Chronos-forecasting not available. Using fallback predictions.")
    CHRONOS_AVAILABLE = False
    BaseChronosPipeline = None

class ModelHandler:
    def __init__(self):
        self.model_name = "amazon/chronos-2" 
        self.pipeline = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.load_model()
    
    def load_model(self):
        """Load Chronos-2 model using the official BaseChronosPipeline"""
        if not CHRONOS_AVAILABLE:
            print("Chronos-forecasting not installed. Using fallback prediction method.")
            return
            
        try:
            print(f"Loading {self.model_name} on {self.device}...")
            
            self.pipeline = BaseChronosPipeline.from_pretrained(
                self.model_name,
                device_map=self.device,
            )
            print("Chronos-2 pipeline loaded successfully.")
            
        except Exception as e:
            print(f"Error loading Chronos-2 model: {e}")
            print("Using fallback prediction method")
            self.pipeline = None
    
    def predict(self, data, horizon=10):
        """Generate predictions using Chronos-2 or fallback."""
        if not CHRONOS_AVAILABLE or self.pipeline is None:
            # Fallback to simple trend-based prediction
            return self._fallback_predict(data, horizon)
            
        try:
            if data is None or not isinstance(data, dict) or 'original' not in data or len(data['original']) < 20:
                return self._fallback_predict(data, horizon)
            
            # --- Chronos-2 Inference ---
            predictions_samples = self.pipeline.predict(
                data['original'],
                prediction_length=horizon,
            )
            
            # Mengambil nilai rata-rata (mean) dari semua sampel atau single trajectory
            if predictions_samples.ndim > 1 and predictions_samples.shape[0] > 1:
                mean_predictions = np.mean(predictions_samples, axis=0)
            elif predictions_samples.ndim > 1 and predictions_samples.shape[0] == 1:
                mean_predictions = predictions_samples[0]
            else:
                mean_predictions = predictions_samples
            
            return mean_predictions
            
        except Exception as e:
            print(f"Prediction error with Chronos: {e}. Using fallback.")
            return self._fallback_predict(data, horizon)
    
    def _fallback_predict(self, data, horizon=10):
        """Fallback prediction method when Chronos is unavailable"""
        try:
            if data is None or not isinstance(data, dict) or 'original' not in data:
                # Return zero predictions if no data
                return np.zeros(horizon)
            
            values = data['original']
            if len(values) < 5:
                return np.zeros(horizon)
            
            # Simple trend extrapolation
            recent_trend = np.polyfit(range(len(values[-20:])), values[-20:], 1)[0]
            predictions = []
            last_value = values[-1]
            
            for i in range(horizon):
                next_value = last_value + recent_trend * (i + 1)
                noise = np.random.normal(0, data.get('std', 1.0) * 0.1)
                predictions.append(next_value + noise)
            
            return np.array(predictions)
            
        except Exception as e:
            print(f"Fallback prediction error: {e}")
            return np.zeros(horizon)