Add integration with AutoTrain

app.py

@@ -1,4 +1,5 @@
 import os
+import uuid
 from pathlib import Path
 
 import streamlit as st
@@ -25,7 +26,7 @@ TASK_TO_ID = {
     "extractive_question_answering": 5,
     "translation": 6,
     "summarization": 8,
-    "single_column_regression": 10,
+    # "single_column_regression": 10,
 }
 
 AUTOTRAIN_TASK_TO_HUB_TASK = {
@@ -36,9 +37,11 @@ AUTOTRAIN_TASK_TO_HUB_TASK = {
     "extractive_question_answering": "question-answering",
     "translation": "translation",
     "summarization": "summarization",
-    "single_column_regression": 10,
+    # "single_column_regression": 10,
 }
 
+HUB_TASK_TO_AUTOTRAIN_TASK = {v: k for k, v in AUTOTRAIN_TASK_TO_HUB_TASK.items()}
+
 ###########
 ### APP ###
 ###########
@@ -63,13 +66,9 @@ if "dataset" in query_params:
 selected_dataset = st.selectbox("Select a dataset", all_datasets, index=all_datasets.index(default_dataset))
 st.experimental_set_query_params(**{"dataset": [selected_dataset]})
 
-# TODO: remove this step once we select real datasets
-# Strip out original dataset name
-# original_dataset_name = dataset_name.split("/")[-1].split("__")[-1]
 
-# In general this will be a list of multiple configs => need to generalise logic here
+# TODO: In general this will be a list of multiple configs => need to generalise logic here
 metadata = get_metadata(selected_dataset)
-print(metadata)
 if metadata is None:
     st.warning("No evaluation metadata found. Please configure the evaluation job below.")
 
@@ -120,16 +119,29 @@ with st.expander("Advanced configuration"):
     # TODO: find a better way to layout these items
     # TODO: propagate this information to payload
     # TODO: make it task specific
+    col_mapping = {}
     with col1:
-        st.markdown("`text` column")
-        st.text("")
-        st.text("")
-        st.text("")
-        st.text("")
-        st.markdown("`target` column")
+        if selected_task == "text-classification":
+            st.markdown("`text` column")
+            st.text("")
+            st.text("")
+            st.text("")
+            st.text("")
+            st.markdown("`target` column")
+        elif selected_task == "question-answering":
+            st.markdown("`context` column")
+            st.text("")
+            st.text("")
+            st.text("")
+            st.text("")
+            st.markdown("`question` column")
     with col2:
-        st.selectbox("This column should contain the text you want to classify", col_names, index=0)
-        st.selectbox("This column should contain the labels you want to assign to the text", col_names, index=1)
+        text_col = st.selectbox("This column should contain the text you want to classify", col_names, index=0)
+        target_col = st.selectbox(
+            "This column should contain the labels you want to assign to the text", col_names, index=1
+        )
+        col_mapping[text_col] = "text"
+        col_mapping[target_col] = "target"
 
 with st.form(key="form"):
 
@@ -138,75 +150,65 @@ with st.form(key="form"):
     selected_models = st.multiselect(
         "Select the models you wish to evaluate", compatible_models
     )  # , compatible_models[0])
-    print(selected_models)
-
     submit_button = st.form_submit_button("Make submission")
 
-    # if submit_button:
-    #     for model in selected_models:
-    #         payload = {
-    #             "username": AUTOTRAIN_USERNAME,
-    #             "task": TASK_TO_ID[metadata[0]["task_id"]],
-    #             "model": model,
-    #             "col_mapping": metadata[0]["col_mapping"],
-    #             "split": selected_split,
-    #             "dataset": original_dataset_name,
-    #             "config": selected_config,
-    #         }
-    # json_resp = http_post(
-    #     path="/evaluate/create", payload=payload, token=HF_TOKEN, domain=AUTOTRAIN_BACKEND_API
-    # ).json()
-    # if json_resp["status"] == 1:
-    #     st.success(f"✅ Successfully submitted model {model} for evaluation with job ID {json_resp['id']}")
-    #     st.markdown(
-    #         f"""
-    #     Evaluation takes appoximately 1 hour to complete, so grab a ☕ or 🍵 while you wait:
-
-    #     * 📊 Click [here](https://huggingface.co/spaces/huggingface/leaderboards) to view the results from your submission
-    #     """
-    #     )
-    # else:
-    #     st.error("🙈 Oh noes, there was an error submitting your submission!")
-
-        # st.write("Creating project!")
-        # payload = {
-        #     "username": AUTOTRAIN_USERNAME,
-        #     "proj_name": "my-eval-project-1",
-        #     "task": TASK_TO_ID[metadata[0]["task_id"]],
-        #     "config": {
-        #         "language": "en",
-        #         "max_models": 5,
-        #         "instance": {
-        #             "provider": "aws",
-        #             "instance_type": "ml.g4dn.4xlarge",
-        #             "max_runtime_seconds": 172800,
-        #             "num_instances": 1,
-        #             "disk_size_gb": 150,
-        #         },
-        #     },
-        # }
-        # json_resp = http_post(
-        #     path="/projects/create", payload=payload, token=HF_TOKEN, domain=AUTOTRAIN_BACKEND_API
-        # ).json()
-        # # print(json_resp)
-
-        # # st.write("Uploading data")
-        # payload = {
-        #     "split": 4,
-        #     "col_mapping": metadata[0]["col_mapping"],
-        #     "load_config": {"max_size_bytes": 0, "shuffle": False},
-        # }
-        # json_resp = http_post(
-        #     path="/projects/522/data/emotion",
-        #     payload=payload,
-        #     token=HF_TOKEN,
-        #     domain=AUTOTRAIN_BACKEND_API,
-        #     params={"type": "dataset", "config_name": "default", "split_name": "train"},
-        # ).json()
-        # print(json_resp)
-
-        # st.write("Training")
-        # json_resp = http_get(
-        #     path="/projects/522/data/start_process", token=HF_TOKEN, domain=AUTOTRAIN_BACKEND_API
-        # ).json()
-        # print(json_resp)
+    if submit_button:
+        project_id = str(uuid.uuid4())[:3]
+        autotrain_task_name = HUB_TASK_TO_AUTOTRAIN_TASK[selected_task]
+        payload = {
+            "username": AUTOTRAIN_USERNAME,
+            "proj_name": f"my-eval-project-{project_id}",
+            "task": TASK_TO_ID[autotrain_task_name],
+            "config": {
+                "language": "en",
+                "max_models": 5,
+                "instance": {
+                    "provider": "aws",
+                    "instance_type": "ml.g4dn.4xlarge",
+                    "max_runtime_seconds": 172800,
+                    "num_instances": 1,
+                    "disk_size_gb": 150,
+                },
+                "evaluation": {
+                    "metrics": [],
+                    "models": selected_models,
+                },
+            },
+        }
+        project_json_resp = http_post(
+            path="/projects/create", payload=payload, token=HF_TOKEN, domain=AUTOTRAIN_BACKEND_API
+        ).json()
+        print(project_json_resp)
+
+        if project_json_resp["created"]:
+            payload = {
+                "split": 4,
+                "col_mapping": col_mapping,
+                "load_config": {"max_size_bytes": 0, "shuffle": False},
+            }
+            data_json_resp = http_post(
+                path=f"/projects/{project_json_resp['id']}/data/{selected_dataset}",
+                payload=payload,
+                token=HF_TOKEN,
+                domain=AUTOTRAIN_BACKEND_API,
+                params={"type": "dataset", "config_name": selected_config, "split_name": selected_split},
+            ).json()
+            print(data_json_resp)
+            if data_json_resp["download_status"] == 1:
+                train_json_resp = http_get(
+                    path=f"/projects/{project_json_resp['id']}/data/start_process",
+                    token=HF_TOKEN,
+                    domain=AUTOTRAIN_BACKEND_API,
+                ).json()
+                print(train_json_resp)
+                if train_json_resp["success"]:
+                    st.success(f"✅ Successfully submitted evaluation job with project ID {project_id}")
+                    st.markdown(
+                        f"""
+                    Evaluation takes appoximately 1 hour to complete, so grab a ☕ or 🍵 while you wait:
+
+                    * 📊 Click [here](https://huggingface.co/spaces/huggingface/leaderboards) to view the results from your submission
+                    """
+                    )
+                else:
+                    st.error("🙈 Oh noes, there was an error submitting your submission!")