add hf token validation

app_text_classification.py

@@ -11,7 +11,14 @@ from text_classification_ui_helpers import (
     try_submit,
     write_column_mapping_to_config,
 )
-from wordings import CONFIRM_MAPPING_DETAILS_MD, INTRODUCTION_MD, USE_INFERENCE_API_TIP, CHECK_LOG_SECTION_RAW
+
+from text_classification import get_example_prediction, HuggingFaceInferenceAPIResponse
+from wordings import (
+  CONFIRM_MAPPING_DETAILS_MD, 
+  INTRODUCTION_MD, 
+  USE_INFERENCE_API_TIP, 
+  CHECK_LOG_SECTION_RAW
+)
 
 MAX_LABELS = 40
 MAX_FEATURES = 20
@@ -225,6 +232,11 @@ def get_demo():
             return gr.update(interactive=False)
         if not column_mapping_accordion.visible:
             return gr.update(interactive=False)
+        _, prediction_response = get_example_prediction(
+            model_id, dataset_id, dataset_config, dataset_split, inference_token
+        )
+        if not isinstance(prediction_response, HuggingFaceInferenceAPIResponse):
+            return gr.update(interactive=False)
         return gr.update(interactive=True)
 
     gr.on(

text_classification.py

@@ -272,7 +272,7 @@ def select_the_first_string_column(ds):
     return None
 
 
-def get_example_prediction(model_id, dataset_id, dataset_config, dataset_split):
+def get_example_prediction(model_id, dataset_id, dataset_config, dataset_split, hf_token):
     # get a sample prediction from the model on the dataset
     prediction_input = None
     prediction_result = None
@@ -284,8 +284,7 @@ def get_example_prediction(model_id, dataset_id, dataset_config, dataset_split):
             prediction_input = ds[0][select_the_first_string_column(ds)]
         else:
             prediction_input = ds[0]["text"]
-        
-        hf_token = os.environ.get(HF_WRITE_TOKEN, default="")
+
         payload = {"inputs": prediction_input, "options": {"use_cache": True}}
         results = hf_inference_api(model_id, hf_token, payload)
 

text_classification_ui_helpers.py

@@ -27,6 +27,7 @@ from wordings import (
     CHECK_LOG_SECTION_RAW,
     get_styled_input,
 )
+import os
 
 MAX_LABELS = 40
 MAX_FEATURES = 20
@@ -214,9 +215,11 @@ def align_columns_and_show_prediction(
     dropdown_placement = [
         gr.Dropdown(visible=False) for _ in range(MAX_LABELS + MAX_FEATURES)
     ]
+    
+    hf_token = os.environ.get("HF_WRITE_TOKEN", default="")
 
     prediction_input, prediction_response = get_example_prediction(
-        model_id, dataset_id, dataset_config, dataset_split
+        model_id, dataset_id, dataset_config, dataset_split, hf_token
     )
 
     if prediction_input is None or prediction_response is None:
@@ -325,7 +328,6 @@ def construct_label_and_feature_mapping(all_mappings, ds_labels, ds_features):
     feature_mapping = all_mappings["features"]
     return label_mapping, feature_mapping
 
-
 def try_submit(m_id, d_id, config, split, inference, inference_token, uid):
     all_mappings = read_column_mapping(uid)
     check_column_mapping_keys_validity(all_mappings)