handle hf api response with custom class

text_classification.py

@@ -7,13 +7,16 @@ import pandas as pd
 from transformers import pipeline
 import requests
 import os
-import time
 
 logger = logging.getLogger(__name__)
 HF_WRITE_TOKEN = "HF_WRITE_TOKEN"
 
 logger = logging.getLogger(__file__)
 
+class HuggingFaceInferenceAPIResponse:
+    def __init__(self, message):
+        self.message = message
+
 
 def get_labels_and_features_from_dataset(ds):
     try:
@@ -286,13 +289,12 @@ def get_example_prediction(model_id, dataset_id, dataset_config, dataset_split):
         payload = {"inputs": prediction_input, "options": {"use_cache": True}}
         results = hf_inference_api(model_id, hf_token, payload)
 
-        if isinstance(results, dict) and "estimated_time" in results.keys():
-            # return the estimated time for the inference api to load
-            # cast the float to int to be concise
-            return prediction_input, str(f"{int(results['estimated_time'])}s")
-
         if isinstance(results, dict) and "error" in results.keys():
-            raise ValueError(results["error"])
+            if "estimated_time" in results.keys():
+                return prediction_input, HuggingFaceInferenceAPIResponse(
+                    f"Estimated time: {int(results['estimated_time'])}s. Please try again later.")
+            return prediction_input, HuggingFaceInferenceAPIResponse(
+                f"Inference Error: {results['error']}.")
         
         while isinstance(results, list):
             if isinstance(results[0], dict):
@@ -303,7 +305,8 @@ def get_example_prediction(model_id, dataset_id, dataset_config, dataset_split):
         }
     except Exception as e:
         # inference api prediction failed, show the error message
-        return prediction_input, e
+        logger.error(f"Get example prediction failed {e}")
+        return prediction_input, None
 
     return prediction_input, prediction_result
 

text_classification_ui_helpers.py

@@ -15,6 +15,7 @@ from text_classification import (
     preload_hf_inference_api,
     get_example_prediction,
     get_labels_and_features_from_dataset,
+    HuggingFaceInferenceAPIResponse,
 )
 from wordings import (
     CHECK_CONFIG_OR_SPLIT_RAW,
@@ -213,24 +214,23 @@ def align_columns_and_show_prediction(
         model_id, dataset_id, dataset_config, dataset_split
     )
 
-    if isinstance(prediction_response, str):
+    if prediction_input is None or prediction_response is None:
         return (
             gr.update(visible=False),
             gr.update(visible=False),
             gr.update(visible=False, open=False),
             gr.update(interactive=False),
-            f"Hugging Face Inference API is loading your model, estimation time {prediction_response}. Please validate again later.",
+            "",
             *dropdown_placement,
         )
 
-    if isinstance(prediction_response, Exception):
-        gr.Warning(f"Inference API loading error: {prediction_response}. Please check your model or Hugging Face token.")
+    if isinstance(prediction_response, HuggingFaceInferenceAPIResponse):
         return (
             gr.update(visible=False),
             gr.update(visible=False),
             gr.update(visible=False, open=False),
             gr.update(interactive=False),
-            "",
+            f"Hugging Face Inference API is loading your model. {prediction_response.message}",
             *dropdown_placement,
         )