No public description

MG_DOCKER_CODES_PIPER_ORIGIN_REV_ID: 784221303

notebooks/community/model_garden/docker_source_codes/model_oss/timesfm/main.py

@@ -4,11 +4,13 @@ import json
 import os
 import flask
 import predictor
+from predictor import PredictionError
 
 # Create the flask app.
 app = flask.Flask(__name__)
 _OK_STATUS = 200
 _INTERNAL_ERROR_STATUS = 500
+_BAD_REQUEST_STATUS = 400
 _HOST = '0.0.0.0'
 
 # Define the predictor and load the checkpoints.
@@ -38,6 +40,18 @@ def predict() -> flask.Response:
         status=_OK_STATUS,
         mimetype='application/json',
     )
+  except PredictionError as e:
+    return flask.Response(
+        json.dumps({'error': str(e)}),
+        status=e.status_code,
+        mimetype='application/json',
+    )
+  except ValueError as e:
+    return flask.Response(
+        json.dumps({'error': str(e)}),
+        status=_BAD_REQUEST_STATUS,
+        mimetype='application/json',
+    )
   except Exception as e:  # pylint: disable=broad-exception-caught
     return flask.Response(
         json.dumps({'error': str(e)}),

notebooks/community/model_garden/docker_source_codes/model_oss/timesfm/predictor.py

@@ -9,9 +9,9 @@ https://www.huggingface.co/google/timesfm-1.0-200m
 
 from collections.abc import Sequence
 import datetime
+import json
 import os
 from typing import Any
-
 import fastapi
 from google.cloud.aiplatform.utils import prediction_utils
 from jax._src import config
@@ -64,11 +64,20 @@ _EXPECTED_FORMAT = """
 """
 
 
+class PredictionError(Exception):
+  """Custom exception for prediction errors."""
+
+  def __init__(self, message: str, status_code: int = _BAD_REQUEST_STATUS):
+    super().__init__(message)
+    self.status_code = status_code
+    self.message = message
+
+
 def _raise_bad_request(message: str):
   message = message + "\n" + _EXPECTED_FORMAT
-  raise HTTPException(
+  raise PredictionError(
+      message=message,
       status_code=_BAD_REQUEST_STATUS,
-      detail=message,
   )
 
 
@@ -476,5 +485,5 @@ class TimesFMPredictor:
         response["timestamp"] = response["timestamp"][: horizon_lens[i]]
 
       predictions.append(response)
-
+    print(f"quantile_forecasts: {json.dumps(quantile_forecasts, indent=2)}")
     return {"predictions": predictions}