Update utils.py

utils.py

@@ -7,6 +7,7 @@ import plotly.graph_objects as go
 import plotly.express as px
 from plotly.subplots import make_subplots
 import spaces
+from chronos import BaseChronosPipeline
 
 def get_indonesian_stocks():
     return {
@@ -213,48 +214,28 @@ def format_large_number(num):
         return f"{num:.2f}"
 
 @spaces.GPU(duration=120)
-def predict_prices(data, model, tokenizer, prediction_days=30):
+def predict_prices(data, model=None, tokenizer=None, prediction_days=30):
     try:
-        prices = data['Close'].values
-        context_length = min(len(prices), 512)
-        input_sequence = prices[-context_length:]
-        price_min = np.min(input_sequence)
-        price_max = np.max(input_sequence)
-        if price_max == price_min:
-            normalized_sequence = np.zeros_like(input_sequence)
-        else:
-            normalized_sequence = (input_sequence - price_min) / (price_max - price_min)
-        VOCAB_SIZE = getattr(model.config, "vocab_size", 2)
-        if VOCAB_SIZE == 2:
-            token_indices = (normalized_sequence > 0.5).astype(np.int64)
-        else:
-            token_indices = (normalized_sequence * (VOCAB_SIZE - 1)).astype(np.int64)
-        prediction_input = torch.tensor(token_indices).unsqueeze(0).to(model.device)
+        prices = data['Close'].values.astype(np.float32)
+        pipeline = BaseChronosPipeline.from_pretrained("amazon/chronos-bolt-base", device_map="auto")
         with torch.no_grad():
-            forecast = model.generate(prediction_input, max_new_tokens=prediction_days, do_sample=True)
-        output_tensor = forecast[0] if isinstance(forecast, (tuple, list)) else forecast
-        predictions_tokens = output_tensor.float().mean(dim=1).squeeze().cpu().numpy()
-        if VOCAB_SIZE == 2:
-            predictions = predictions_tokens * (price_max - price_min) + price_min
-        else:
-            predictions = (predictions_tokens / (VOCAB_SIZE - 1)) * (price_max - price_min) + price_min
-        if predictions.ndim == 0:
-            predictions = np.array([predictions.item()])
-        pred_len = len(predictions)
+            forecast = pipeline.predict(context=torch.tensor(prices), prediction_length=prediction_days)
+        mean_forecast = forecast.mean(dim=1).squeeze().cpu().numpy()
+        pred_len = len(mean_forecast)
         last_price = prices[-1]
-        predicted_high = np.max(predictions)
-        predicted_low = np.min(predictions)
-        predicted_mean = np.mean(predictions)
+        predicted_high = np.max(mean_forecast)
+        predicted_low = np.min(mean_forecast)
+        predicted_mean = np.mean(mean_forecast)
         change_pct = ((predicted_mean - last_price) / last_price) * 100
         return {
-            'values': predictions,
+            'values': mean_forecast,
             'dates': pd.date_range(start=data.index[-1] + timedelta(days=1), periods=pred_len, freq='D'),
             'high_30d': predicted_high,
             'low_30d': predicted_low,
             'mean_30d': predicted_mean,
             'change_pct': change_pct,
             'summary': f"""
-AI Model: Amazon Chronos-Bolt (Binary Quantization)
+AI Model: Amazon Chronos-Bolt (Base)
 Prediction Period: {pred_len} days
 Expected Change: {change_pct:.2f}%
 Confidence: Medium