Add connection to AutoTrain

app.py

@@ -1,41 +1,57 @@
-import streamlit as st
-from huggingface_hub import DatasetFilter, HfApi, ModelFilter
-
-api = HfApi()
+import os
+from pathlib import Path
 
+import streamlit as st
+from dotenv import load_dotenv
 
-def get_metadata(dataset_name):
-    filt = DatasetFilter(dataset_name=dataset_name)
-    data = api.list_datasets(filter=filt, full=True)
-    return data[0].cardData["train-eval-index"]
+from utils import get_compatible_models, get_metadata, http_post
 
+if Path(".env").is_file():
+    load_dotenv(".env")
 
-def get_compatible_models(task, dataset_name):
-    filt = ModelFilter(task=task, trained_dataset=dataset_name)
-    compatible_models = api.list_models(filter=filt)
-    return [model.modelId for model in compatible_models]
+HF_TOKEN = os.getenv("HF_TOKEN")
+AUTOTRAIN_USERNAME = os.getenv("AUTOTRAIN_USERNAME")
+AUTOTRAIN_BACKEND_API = os.getenv("AUTOTRAIN_BACKEND_API")
 
 
 with st.form(key="form"):
 
-    dataset_name = st.selectbox("Select a dataset to evaluate on", ["lewtun/autoevaluate_emotion"])
+    dataset_name = st.selectbox("Select a dataset to evaluate on", ["lewtun/autoevaluate__emotion"])
 
+    # 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
     metadata = get_metadata(dataset_name)
-    # st.write(metadata)
 
     dataset_config = st.selectbox("Select the subset to evaluate on", [metadata[0]["config"]])
 
     splits = metadata[0]["splits"]
+    split_names = list(splits.values())
+    eval_split = splits.get("eval_split", split_names[0])
 
-    # st.write(splits)
-
-    evaluation_split = st.selectbox("Select the split to evaluate on", [v for d in splits for k, v in d.items()])
+    selected_split = st.selectbox("Select the split to evaluate on", split_names, index=split_names.index(eval_split))
 
-    compatible_models = get_compatible_models(metadata[0]["task"], dataset_name.split("/")[-1].split("_")[-1])
+    compatible_models = get_compatible_models(metadata[0]["task"], original_dataset_name)
 
-    options = st.multiselect("Select the models you wish to evaluate", compatible_models, compatible_models[0])
+    selected_models = st.multiselect("Select the models you wish to evaluate", compatible_models, compatible_models[0])
 
     submit_button = st.form_submit_button("Make Submission")
 
     if submit_button:
-        st.success(f"✅ Evaluation was successfully submitted for evaluation with job ID 42")
+        for model in selected_models:
+            payload = {
+                "username": AUTOTRAIN_USERNAME,
+                "task": 1,
+                "model": model,
+                "col_mapping": {"sentence": "text", "label": "target"},
+                "split": selected_split,
+                "dataset": original_dataset_name,
+                "config": dataset_config,
+            }
+            json_resp = http_post(
+                path="/evaluate/create", payload=payload, token=HF_TOKEN, domain=AUTOTRAIN_BACKEND_API
+            ).json()
+
+            st.success(f"✅ Successfully submitted model {model} for evaluation with job ID {json_resp['id']}")

requirements.txt

@@ -1 +1,2 @@
-huggingface-hub==0.4.0
+huggingface-hub==0.4.0
+python-dotenv

utils.py

@@ -0,0 +1,37 @@
+import requests
+from huggingface_hub import DatasetFilter, HfApi, ModelFilter
+
+api = HfApi()
+
+
+def get_auth_headers(token: str, prefix: str = "autonlp"):
+    return {"Authorization": f"{prefix} {token}"}
+
+
+def http_post(
+    path: str,
+    token: str,
+    payload=None,
+    domain: str = None,
+) -> requests.Response:
+    """HTTP POST request to the AutoNLP API, raises UnreachableAPIError if the API cannot be reached"""
+    try:
+        response = requests.post(
+            url=domain + path, json=payload, headers=get_auth_headers(token=token), allow_redirects=True
+        )
+    except requests.exceptions.ConnectionError:
+        print("❌ Failed to reach AutoNLP API, check your internet connection")
+    response.raise_for_status()
+    return response
+
+
+def get_metadata(dataset_name):
+    filt = DatasetFilter(dataset_name=dataset_name)
+    data = api.list_datasets(filter=filt, full=True)
+    return data[0].cardData["train-eval-index"]
+
+
+def get_compatible_models(task, dataset_name):
+    filt = ModelFilter(task=task, trained_dataset=dataset_name)
+    compatible_models = api.list_models(filter=filt)
+    return [model.modelId for model in compatible_models]