Fix the first submission for the task case and tidy up code a bit

src/submission_uploader.py

@@ -3,7 +3,7 @@ import logging
 import os
 import time
 from tempfile import TemporaryDirectory
-from typing import Dict, List, Optional
+from typing import List, Optional
 
 import jsonlines
 from huggingface_hub import CommitOperationAdd  # type: ignore[import]
@@ -30,13 +30,13 @@ class SubmissionUploader:
     def __init__(self, dataset_id: str, private_dataset_id: str):
         self._api = HfApi(token=os.environ["HF_TOKEN"])
         self._fs = HfFileSystem(token=os.environ["HF_TOKEN"])
-        self._dataset_id = dataset_id
-        self._private_dataset_id = private_dataset_id
+        self._results_dataset_id = dataset_id
+        self._requests_dataset_id = private_dataset_id
 
     def _get_previous_pr(self, pr_title: str) -> Optional[Discussion]:
-        """Searches among discussions of dataset repo for a PR with the given title."""
+        """Searches among discussions of the results dataset for a PR with the given title."""
         try:
-            discussions = self._api.get_repo_discussions(repo_id=self._dataset_id, repo_type="dataset")
+            discussions = self._api.get_repo_discussions(repo_id=self._results_dataset_id, repo_type="dataset")
         except Exception:
             return None
         for discussion in discussions:
@@ -44,22 +44,6 @@ class SubmissionUploader:
                 return discussion
         return None
 
-    def _get_metadata(
-        self,
-        model_name_pretty: str,
-        model_availability: str,
-        urls: Optional[str],
-        context_size: str,
-        submitted_by: str,
-    ) -> Dict[str, Optional[str]]:
-        return {
-            "model_name": model_name_pretty,
-            "model_availability": model_availability,
-            "urls": urls,
-            "context_size": context_size,
-            "submitted_by": submitted_by,
-        }
-
     def _upload_request(
         self,
         task_id: str,
@@ -74,6 +58,7 @@ class SubmissionUploader:
         pr_url: str,
         temp_directory: str,
     ) -> List[CommitOperationAdd]:
+        """Adds a file with metadata about the current request to the requests dataset."""
         request_metadata = {
             "model_folder": model_folder,
             "model_name_pretty": model_name_pretty,
@@ -90,7 +75,11 @@ class SubmissionUploader:
         with open(os.path.join(temp_directory, "request_metadata.json"), "w") as f:
             json.dump(request_metadata, f)
 
-        num_requests_already_present = len(self._fs.ls(f"datasets/{self._private_dataset_id}/{task_id}/"))
+        num_requests_already_present = (
+            len(self._fs.ls(f"datasets/{self._requests_dataset_id}/{task_id}/"))
+            if self._fs.isdir(f"datasets/{self._requests_dataset_id}/{task_id}/")
+            else 0
+        )
         commit_operations = [
             CommitOperationAdd(
                 path_in_repo=f"{task_id}/{num_requests_already_present}_{model_folder}.json",
@@ -105,6 +94,7 @@ class SubmissionUploader:
         model_folder: str,
         filenames: List[str],
     ) -> List[CommitOperationAdd]:
+        """Adds all files with current model's predictions to the results dataset."""
         commit_operations = [
             CommitOperationAdd(
                 path_in_repo=f"{task_id}/predictions/{model_folder}/{os.path.basename(filename)}",
@@ -115,6 +105,7 @@ class SubmissionUploader:
         return commit_operations
 
     def _compute_metrics_for_predictions(self, task_id: str, filenames: List[str], temp_directory: str) -> None:
+        """Computes metrics for each submitted file with the current model's predictions."""
         metrics_module = METRICS[task_id]
         assert metrics_module is not None, f"Computing metrics for {task_id} is not supported."
         metrics_module.reset()
@@ -153,18 +144,20 @@ class SubmissionUploader:
         submitted_by: str,
         temp_directory: str,
     ) -> List[CommitOperationAdd]:
+        """Adds files with the current model's metrics values to the results dataset."""
         final_results = {}
         with open(os.path.join(temp_directory, "final_metrics.json"), "r") as f:
             metrics = json.load(f)
         final_results.update(metrics)
-        metadata_dict = self._get_metadata(
-            model_name_pretty=model_name_pretty,
-            model_availability=model_availability,
-            urls=urls,
-            context_size=context_size,
-            submitted_by=submitted_by,
+        final_results.update(
+            {
+                "model_name": model_name_pretty,
+                "model_availability": model_availability,
+                "urls": urls,
+                "context_size": context_size,
+                "submitted_by": submitted_by,
+            }
         )
-        final_results.update(metadata_dict)
 
         with jsonlines.open(os.path.join(temp_directory, "final_results.jsonl"), "w") as writer:
             writer.write(final_results)
@@ -189,6 +182,7 @@ class SubmissionUploader:
         comment: Optional[str],
         filenames: Optional[List[str]],
     ):
+        """Verifies that all necessary arguments are not None (and also runs other sanity checks)."""
         assert task_pretty and task_pretty in TASKS_PRETTY_REVERSE, "Please, select one of the supported tasks."
         assert model_folder, "Please, specify non-empty name for a directory with a model's results."
         assert model_name_pretty, "Please, specify non-empty name for a model."
@@ -238,15 +232,17 @@ class SubmissionUploader:
 
             logging.info("Checking if this request has already been submitted...")
             if not force:
-                if model_folder in self._fs.ls(f"datasets/{self._dataset_id}/{task_id}/predictions"):
+                if self._fs.isdir(f"datasets/{self._results_dataset_id}/{task_id}/predictions/{model_folder}"):
                     return styled_warning(
-                        f"{model_folder} is already present in {self._dataset_id}, please, select another folder name."
+                        f"{model_folder} is already present in {self._results_dataset_id}, please, select another folder name."
                     )
 
                 prev_pr = self._get_previous_pr(pr_title)
                 if prev_pr is not None:
-                    url = f"https://huggingface.co/datasets/{self._dataset_id}/discussions/{prev_pr.num}"
-                    return styled_warning(f"{self._dataset_id} already has an open PR for this submission: {url}.")
+                    url = f"https://huggingface.co/datasets/{self._results_dataset_id}/discussions/{prev_pr.num}"
+                    return styled_warning(
+                        f"{self._results_dataset_id} already has an open PR for this submission: {url}."
+                    )
 
             logging.info("Processing predictions...")
             predictions_commit_operations = self._upload_predictions(
@@ -271,9 +267,9 @@ class SubmissionUploader:
                     temp_directory=str(d),
                 )
 
-                logging.info(f"Creating commit to results dataset...")
+                logging.info("Creating commit to the results dataset...")
                 new_pr = self._api.create_commit(
-                    repo_id=self._dataset_id,
+                    repo_id=self._results_dataset_id,
                     operations=predictions_commit_operations + results_commit_operations,
                     commit_message=pr_title,
                     commit_description=f"""New submission to {task_pretty} task in 🏟️ Long Code Arena benchmark!\n* Model name: {model_name_pretty}\n* Model availability: {model_availability}\n* Context Size: {context_size}\n* Relevant URLs: {urls}\n* Submitted By: {submitted_by}""",
@@ -281,7 +277,7 @@ class SubmissionUploader:
                     repo_type="dataset",
                 )
 
-                logging.info(f"Creating commit to requests dataset...")
+                logging.info("Creating commit to the requests dataset...")
                 request_commit_operations = self._upload_request(
                     task_id=task_id,
                     model_folder=model_folder,
@@ -296,7 +292,7 @@ class SubmissionUploader:
                     pr_url=new_pr.pr_url,
                 )
                 self._api.create_commit(
-                    repo_id=self._private_dataset_id,
+                    repo_id=self._requests_dataset_id,
                     operations=request_commit_operations,
                     commit_message=pr_title,
                     commit_description=f"""New submission to {task_pretty} task in 🏟️ Long Code Arena benchmark!\n* Model name: {model_name_pretty}\n* Model availability: {model_availability}\n* Context Size: {context_size}\n* Relevant URLs: {urls}\n* Submitted By: {submitted_by}\n* PR: {new_pr.pr_url}\n* Contact information: {contact_information}\n* Comment: {comment}""",
@@ -307,7 +303,6 @@ class SubmissionUploader:
                 return styled_message(f"🎉 PR created at {new_pr.pr_url}.")
 
         except Exception as e:
-            logging.exception(e)
             exception_msg = str(e)
             if exception_msg and os.environ["PRIVATE_DATASET_ID"] in exception_msg:
                 exception_msg = exception_msg.replace(os.environ["PRIVATE_DATASET_ID"], "{private_dataset}")