pieces

dependencies/requirements_0.txt

@@ -1,3 +1,4 @@
 prophet==1.1.5
 pandas==2.1.3
-plotly==5.18.0
+plotly==5.18.0
+yfinance==0.2.38

pieces/GetYahooFinanceDataPiece/metadata.json

@@ -0,0 +1,24 @@
+{
+    "name": "GetYahooFinanceDataPiece",
+    "description": "This piece gets data from Yahoo Finance.",
+    "dependency": {
+        "requirements_file": "requirements_0.txt"
+    },
+    "container_resources": {
+        "requests": {
+            "cpu": "100m",
+            "memory": "128Mi"
+        },
+        "limits": {
+            "cpu": "500m",
+            "memory": "512Mi"
+        }
+    },
+    "tags": [
+        "Example"
+    ],
+    "style": {
+        "node_label": "Get Yahoo Finance Data",
+        "icon_class_name": "fa-solid:database"
+    }
+}

pieces/GetYahooFinanceDataPiece/models.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel, Field
+from datetime import date
+
+
+
+class InputModel(BaseModel):
+    ticker: str = Field(
+        description="Ticker of the stock to get data from."
+    ) # TODO change to ENUM ?
+    start_date: date = Field(
+        description="Start date of the data to get."
+    )
+    end_date: date = Field(
+        description="End date of the data to get."
+    )
+
+class OutputModel(BaseModel):
+    data_path: str = Field(
+        description="Path to the file containing the trained model."
+    )

pieces/GetYahooFinanceDataPiece/piece.py

@@ -0,0 +1,24 @@
+from domino.base_piece import BasePiece
+from .models import InputModel, OutputModel
+import yfinance as yf
+from pathlib import Path
+
+
+class GetYahooFinanceDataPiece(BasePiece):
+    """
+    This Piece trains a Prophet model using the data provided in the input file.
+    """
+    def piece_function(self, input_data: InputModel):
+        ticker = input_data.ticker
+        start_date = input_data.start_date
+        end_date = input_data.end_date
+
+        df = yf.download(ticker, start=start_date, end=end_date)
+        df.reset_index(inplace=True)
+
+        df_path = Path(self.results_path) / f"{ticker}_data.csv"
+        df.to_csv(df_path, index=False)
+
+        return OutputModel(
+            data_path=str(df_path)
+        )

pieces/ProphetPredictPiece/metadata.json

@@ -0,0 +1,24 @@
+{
+    "name": "ProphetPredictPiece",
+    "description": "Piece to predict using Prophet",
+    "dependency": {
+        "requirements_file": "requirements_0.txt"
+    },
+    "container_resources": {
+        "requests": {
+            "cpu": "100m",
+            "memory": "128Mi"
+        },
+        "limits": {
+            "cpu": "500m",
+            "memory": "512Mi"
+        }
+    },
+    "tags": [
+        "Prophet"
+    ],
+    "style": {
+        "node_label": "Prophet Predict",
+        "icon_class_name": "fa-solid:database"
+    }
+}

pieces/ProphetPredictPiece/models.py

@@ -0,0 +1,26 @@
+from pydantic import BaseModel, Field
+from enum import Enum
+from typing import List
+from datetime import date
+
+
+class InputModel(BaseModel):
+    model_path: str = Field(
+        title="Model Path",
+        description="Path to the file containing the trained model."
+    )
+    periods: int = Field(
+        title="Periods",
+        description="Number of periods to forecast."
+    )
+
+
+class OutputModel(BaseModel):
+    forecast_data_path: str = Field(
+        title="Forecast Data Path",
+        description="Path to the file containing the forecast data."
+    )
+    forecast_figure_path: str = Field(
+        title="Forecast Figure Path",
+        description="Path to the file containing the results figure."
+    )

pieces/ProphetPredictPiece/piece.py

@@ -0,0 +1,40 @@
+from domino.base_piece import BasePiece
+from .models import InputModel, OutputModel
+import pandas as pd
+from prophet import Prophet
+import pickle
+from pathlib import Path
+from prophet.plot import plot_plotly, plot_components_plotly
+
+
+
+class ProphetPredictPiece(BasePiece):
+    """
+    This Piece uses a trained Prophet model to make predictions on new data.
+    """
+    def piece_function(self, input_data: InputModel):
+
+        with open(input_data.model_path, "rb") as f:
+            model = pickle.load(f)
+
+        future = model.make_future_dataframe(periods=input_data.periods)
+        forecast = model.predict(future)
+
+        self.results_path = Path(self.results_path)
+
+        forecast_data_path = self.results_path / "forecast_data.csv"
+        forecast.to_csv(forecast_data_path, index=False)
+
+        forecast_figure_path = self.results_path / "forecast_figure.json"
+        forecast_figure = plot_plotly(model, forecast)
+
+        forecast_figure.write_json(str(forecast_figure_path))
+        self.display_result = {
+            "file_type": "plotly_json",
+            "file_path": str(forecast_figure_path)
+        }
+
+        return OutputModel(
+            forecast_data_path=str(forecast_data_path),
+            forecast_figure_path=str(forecast_figure_path),    
+        )

pieces/ProphetPredictPiece/test_example_simple_piece.py

@@ -0,0 +1,15 @@
+from domino.testing import piece_dry_run
+
+def test_example_simple_piece():
+    input_data = dict(
+        distribution_name="gaussian",
+        distribution_mean=0.,
+        distribution_sd=1.
+    )
+    output_data = piece_dry_run(
+        "ExampleSimplePiece",
+        input_data
+    )
+
+    assert output_data["message"] is not None
+    assert output_data["sample_result"] is not None

pieces/ProphetTrainModelPiece/metadata.json

@@ -1,6 +1,6 @@
 {
-    "name": "ExampleSimplePiece",
-    "description": "This is an example of a simple Domino Piece",
+    "name": "ProphetTrainModelPiece",
+    "description": "Piece to train a prophet model",
     "dependency": {
         "requirements_file": "requirements_0.txt"
     },
@@ -15,10 +15,10 @@
         }
     },
     "tags": [
-        "Example"
+        "Prophet"
     ],
     "style": {
-        "node_label": "Simple Piece",
+        "node_label": "Prophet Train",
         "icon_class_name": "fa-solid:database"
     }
 }

pieces/ProphetTrainModelPiece/models.py

@@ -17,14 +17,15 @@ class SeasonalityMode(str, Enum):
 
 class InputModel(BaseModel):
     input_data_file: str = Field(
-        description="Path to the input data file. Accepted formats: `.csv`, `.json`."
-    )
-    datetime_column_name: str = Field(
-        description="Name of the column containing the datetime values."
-    )
-    target_column_name: str = Field(
-        description="Name of the column containing the target values."
-    )
+        title="Input Data File",
+        description="Path to the input data file. Accepted formats: `.csv`, `.json`. Should use the following format: `ds` (datetime), `y` (target).",
+    )
+    # datetime_column_name: str = Field(
+    #     description="Name of the column containing the datetime values."
+    # )
+    # target_column_name: str = Field(
+    #     description="Name of the column containing the target values."
+    # )
     test_set_percentage: float = Field(
         default=20.0,
         ge=1,
@@ -55,6 +56,6 @@ class OutputModel(BaseModel):
     model_file_path: str = Field(
         description="Path to the file containing the trained model."
     )
-    results_figure_file_path: str = Field(
-        description="Path to the file containing the results figure."
-    )
+    # results_figure_file_path: str = Field(
+    #     description="Path to the file containing the results figure."
+    # )

pieces/ProphetTrainModelPiece/piece.py

@@ -2,9 +2,8 @@
 from .models import InputModel, OutputModel
 import pandas as pd
 from prophet import Prophet
-from prophet.serialize import model_to_json
-import plotly.graph_objs as go
-
+import pickle
+from pathlib import Path
 
 class ProphetTrainModelPiece(BasePiece):
     """
@@ -21,9 +20,14 @@ def piece_function(self, input_data: InputModel):
         else:
             raise ValueError("File format not supported. Please pass a CSV or JSON file.")
 
+        model = Prophet()
+        model.fit(df)
 
+        # Serialize model
+        model_file_path = self.results_path / "prophet_model.json"
+        with open(str(model_file_path), "wb") as f:
+            pickle.dump(model, f)
 
         return OutputModel(
-            model_file_path=,
-            results_figure_file_path=
+            model_file_path=str(model_file_path),
         )