Spaces:
Sleeping
Sleeping
import gradio as gr | |
from wordings import CONFIRM_MAPPING_DETAILS_FAIL_RAW | |
import json | |
import os | |
import logging | |
import uuid | |
import threading | |
from io_utils import read_column_mapping, write_column_mapping, save_job_to_pipe, write_log_to_user_file | |
import datasets | |
import collections | |
from text_classification import get_labels_and_features_from_dataset, check_model, get_example_prediction | |
from transformers.pipelines import TextClassificationPipeline | |
MAX_LABELS = 20 | |
MAX_FEATURES = 20 | |
HF_REPO_ID = 'HF_REPO_ID' | |
HF_SPACE_ID = 'SPACE_ID' | |
HF_WRITE_TOKEN = 'HF_WRITE_TOKEN' | |
CONFIG_PATH = "./config.yaml" | |
def check_dataset_and_get_config(dataset_id): | |
try: | |
write_column_mapping(None) | |
configs = datasets.get_dataset_config_names(dataset_id) | |
return gr.Dropdown(configs, value=configs[0], visible=True) | |
except Exception: | |
# Dataset may not exist | |
pass | |
def check_dataset_and_get_split(dataset_id, dataset_config): | |
try: | |
splits = list(datasets.load_dataset(dataset_id, dataset_config).keys()) | |
return gr.Dropdown(splits, value=splits[0], visible=True) | |
except Exception: | |
# Dataset may not exist | |
# gr.Warning(f"Failed to load dataset {dataset_id} with config {dataset_config}: {e}") | |
pass | |
def write_column_mapping_to_config(dataset_id, dataset_config, dataset_split, *labels): | |
ds_labels, ds_features = get_labels_and_features_from_dataset(dataset_id, dataset_config, dataset_split) | |
if labels is None: | |
return | |
labels = [*labels] | |
all_mappings = read_column_mapping(CONFIG_PATH) | |
if "labels" not in all_mappings.keys(): | |
all_mappings["labels"] = dict() | |
for i, label in enumerate(labels[:MAX_LABELS]): | |
if label: | |
all_mappings["labels"][label] = ds_labels[i] | |
if "features" not in all_mappings.keys(): | |
all_mappings["features"] = dict() | |
for i, feat in enumerate(labels[MAX_LABELS:(MAX_LABELS + MAX_FEATURES)]): | |
if feat: | |
all_mappings["features"][feat] = ds_features[i] | |
write_column_mapping(all_mappings) | |
def list_labels_and_features_from_dataset(ds_labels, ds_features, model_id2label): | |
model_labels = list(model_id2label.values()) | |
lables = [gr.Dropdown(label=f"{label}", choices=model_labels, value=model_id2label[i], interactive=True, visible=True) for i, label in enumerate(ds_labels[:MAX_LABELS])] | |
lables += [gr.Dropdown(visible=False) for _ in range(MAX_LABELS - len(lables))] | |
# TODO: Substitute 'text' with more features for zero-shot | |
features = [gr.Dropdown(label=f"{feature}", choices=ds_features, value=ds_features[0], interactive=True, visible=True) for feature in ['text']] | |
features += [gr.Dropdown(visible=False) for _ in range(MAX_FEATURES - len(features))] | |
return lables + features | |
def check_model_and_show_prediction(model_id, dataset_id, dataset_config, dataset_split): | |
ppl = check_model(model_id) | |
if ppl is None or not isinstance(ppl, TextClassificationPipeline): | |
gr.Warning("Please check your model.") | |
return ( | |
gr.update(visible=False), | |
gr.update(visible=False), | |
*[gr.update(visible=False) for _ in range(MAX_LABELS + MAX_FEATURES)] | |
) | |
dropdown_placement = [gr.Dropdown(visible=False) for _ in range(MAX_LABELS + MAX_FEATURES)] | |
if ppl is None: # pipeline not found | |
gr.Warning("Model not found") | |
return ( | |
gr.update(visible=False), | |
gr.update(visible=False), | |
gr.update(visible=False, open=False), | |
*dropdown_placement | |
) | |
model_id2label = ppl.model.config.id2label | |
ds_labels, ds_features = get_labels_and_features_from_dataset(dataset_id, dataset_config, dataset_split) | |
# when dataset does not have labels or features | |
if not isinstance(ds_labels, list) or not isinstance(ds_features, list): | |
gr.Warning(CONFIRM_MAPPING_DETAILS_FAIL_RAW) | |
return ( | |
gr.update(visible=False), | |
gr.update(visible=False), | |
gr.update(visible=False, open=False), | |
*dropdown_placement | |
) | |
column_mappings = list_labels_and_features_from_dataset( | |
ds_labels, | |
ds_features, | |
model_id2label, | |
) | |
# when labels or features are not aligned | |
# show manually column mapping | |
if collections.Counter(model_id2label.values()) != collections.Counter(ds_labels) or ds_features[0] != 'text': | |
gr.Warning(CONFIRM_MAPPING_DETAILS_FAIL_RAW) | |
return ( | |
gr.update(visible=False), | |
gr.update(visible=False), | |
gr.update(visible=True, open=True), | |
*column_mappings | |
) | |
prediction_input, prediction_output = get_example_prediction(ppl, dataset_id, dataset_config, dataset_split) | |
return ( | |
gr.update(value=prediction_input, visible=True), | |
gr.update(value=prediction_output, visible=True), | |
gr.update(visible=True, open=False), | |
*column_mappings | |
) | |
def get_logs_file(uid): | |
file = open(f"./tmp/{uid}_log") | |
contents = file.readlines() | |
file.close() | |
return '\n'.join(contents) | |
def try_submit(m_id, d_id, config, split, local): | |
all_mappings = read_column_mapping(CONFIG_PATH) | |
if all_mappings is None: | |
gr.Warning(CONFIRM_MAPPING_DETAILS_FAIL_RAW) | |
return gr.update(interactive=True) | |
if "labels" not in all_mappings.keys(): | |
gr.Warning(CONFIRM_MAPPING_DETAILS_FAIL_RAW) | |
return gr.update(interactive=True) | |
label_mapping = all_mappings["labels"] | |
if "features" not in all_mappings.keys(): | |
gr.Warning(CONFIRM_MAPPING_DETAILS_FAIL_RAW) | |
return gr.update(interactive=True) | |
feature_mapping = all_mappings["features"] | |
# TODO: Set column mapping for some dataset such as `amazon_polarity` | |
if local: | |
command = [ | |
"python", | |
"cli.py", | |
"--loader", "huggingface", | |
"--model", m_id, | |
"--dataset", d_id, | |
"--dataset_config", config, | |
"--dataset_split", split, | |
"--hf_token", os.environ.get(HF_WRITE_TOKEN), | |
"--discussion_repo", os.environ.get(HF_REPO_ID) or os.environ.get(HF_SPACE_ID), | |
"--output_format", "markdown", | |
"--output_portal", "huggingface", | |
"--feature_mapping", json.dumps(feature_mapping), | |
"--label_mapping", json.dumps(label_mapping), | |
"--scan_config", "../config.yaml", | |
] | |
eval_str = f"[{m_id}]<{d_id}({config}, {split} set)>" | |
logging.info(f"Start local evaluation on {eval_str}") | |
uid = uuid.uuid4() | |
save_job_to_pipe(uid, command, threading.Lock()) | |
write_log_to_user_file(uid, f"Start local evaluation on {eval_str}. Please wait for your job to start...\n") | |
gr.Info(f"Start local evaluation on {eval_str}") | |
return ( | |
gr.update(interactive=False), | |
gr.update(value=get_logs_file(uid), visible=True, interactive=False, every=5)) | |
else: | |
gr.Info("TODO: Submit task to an endpoint") | |
return (gr.update(interactive=True), # Submit button | |
gr.update(visible=False)) | |