Quick fix for inference API not ready

text_classification.py

@@ -7,9 +7,13 @@ import pandas as pd
 from transformers import pipeline
 import requests
 import os
+import time
 
 HF_WRITE_TOKEN = "HF_WRITE_TOKEN"
 
+logger = logging.getLogger(__file__)
+
+
 def get_labels_and_features_from_dataset(ds):
     try:
         dataset_features = ds.features
@@ -72,13 +76,19 @@ def hf_inference_api(model_id, hf_token, payload):
     )
     url = f"{hf_inference_api_endpoint}/models/{model_id}"
     headers = {"Authorization": f"Bearer {hf_token}"}
-    response = requests.post(url, headers=headers, json=payload)
-    if response.status_code != 200:
-        logging.ERROR(f"Request to inference API returns {response.status_code}")
-    try:
-        return response.json()
-    except Exception:
-        return {"error": response.content}
+    output = {"error": "First attemp"}
+    attempt = 30
+    while "error" in output and attempt > 0:
+        response = requests.post(url, headers=headers, json=payload)
+        if response.status_code != 200:
+            logging.error(f"Request to inference API returns {response.status_code}")
+        try:
+            return response.json()
+        except Exception:
+            logging.error(f"{response.content}")
+            output = {"error": response.content}
+        attempt -= 1
+        time.sleep(2)
 
 def check_model_pipeline(model_id):
     try:
@@ -262,7 +272,7 @@ def get_example_prediction(model_id, dataset_id, dataset_config, dataset_split):
         ds = datasets.load_dataset(dataset_id, dataset_config)[dataset_split]
         if "text" not in ds.features.keys():
             # Dataset does not have text column
-            prediction_input = ds[0][ds.features.keys()[0]]
+            prediction_input = ds[0][list(ds.features.keys())[0]]
         else:
             prediction_input = ds[0]["text"]
 
@@ -276,8 +286,9 @@ def get_example_prediction(model_id, dataset_id, dataset_config, dataset_split):
         prediction_result = {
             f'{result["label"]}': result["score"] for result in results
         }
-    except Exception:
+    except Exception as e:
         # Pipeline prediction failed, need to provide labels
+        logger.warn(f"Pipeline prediction failed due to {e}")
         return prediction_input, None
 
     return prediction_input, prediction_result

text_classification_ui_helpers.py

@@ -209,6 +209,7 @@ def align_columns_and_show_prediction(
         gr.Dropdown(visible=False) for _ in range(MAX_LABELS + MAX_FEATURES)
     ]
 
+    # FIXME: prefiction_output could be None
     prediction_input, prediction_output = get_example_prediction(
         model_id, dataset_id, dataset_config, dataset_split
     )