add distribution type

app.py

@@ -1,8 +1,14 @@
 import gradio as gr
 import pandas as pd
+
 from gluonts.dataset.pandas import PandasDataset
 from gluonts.dataset.split import split
 from gluonts.torch.model.deepar import DeepAREstimator
+from gluonts.torch.distributions import (
+    NegativeBinomialOutput,
+    StudentTOutput,
+    NormalOutput,
+)
 from gluonts.evaluation import Evaluator, make_evaluation_predictions
 
 from make_plot import plot_forecast, plot_train_test
@@ -32,6 +38,7 @@ def train_and_forecast(
     prediction_length,
     rolling_windows,
     epochs,
+    distribution,
     progress=gr.Progress(track_tqdm=True),
 ):
     if not input_data:
@@ -54,7 +61,14 @@ def train_and_forecast(
 
     training_data, test_gen = split(gluon_df, offset=row_offset)
 
+    if distribution == "StudentT":
+        distr_output = StudentTOutput()
+    elif distribution == "Normal":
+        distr_output = NormalOutput()
+    else:
+        distr_output = NegativeBinomialOutput()
     estimator = DeepAREstimator(
+        distr_output=distr_output,
         prediction_length=prediction_length,
         freq=gluon_df.freq,
         trainer_kwargs=dict(max_epochs=epochs),
@@ -108,6 +122,11 @@ with gr.Blocks() as demo:
             )
             windows = gr.Number(value=3, label="Number of Windows", precision=0)
             epochs = gr.Number(value=10, label="Number of Epochs", precision=0)
+            distribution = gr.Radio(
+                choices=["StudentT", "Negative Binomial", "Normal"],
+                value="StudentT",
+                label="Distribution",
+            )
 
     with gr.Row(label="Dataset"):
         upload_btn = gr.UploadButton(label="Upload")
@@ -122,7 +141,7 @@ with gr.Blocks() as demo:
     )
     train_btn.click(
         fn=train_and_forecast,
-        inputs=[upload_btn, prediction_length, windows, epochs],
+        inputs=[upload_btn, prediction_length, windows, epochs, distribution],
         outputs=[plot, json],
     )