import json import logging import os from tempfile import TemporaryDirectory from typing import Dict, List, Optional import jsonlines from huggingface_hub import CommitOperationAdd # type: ignore[import] from huggingface_hub import Discussion, HfApi, HfFileSystem from tqdm import tqdm from .evaluation import METRICS from .formatting import styled_error, styled_message, styled_warning from .tasks import TASKS_PRETTY_REVERSE class AlreadyExists(Exception): pass class SubmissionUploader: """Class for adding new files to a dataset on a Hub and opening a PR. Heavily influenced by these amazing spaces: * https://huggingface.co/spaces/safetensors/convert * https://huggingface.co/spaces/gaia-benchmark/leaderboard """ def __init__(self, dataset_id: str): self._api = HfApi(token=os.environ["HF_TOKEN"]) self._fs = HfFileSystem(token=os.environ["HF_TOKEN"]) self._dataset_id = 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.""" try: discussions = self._api.get_repo_discussions( repo_id=self._dataset_id, repo_type="dataset" ) except Exception: return None for discussion in discussions: if ( discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title ): return discussion return None def _get_metadata( self, model_name_pretty: str, model_availability: str, urls: str, context_size: str, submitted_by: str, ) -> Dict[str, str]: return { "model_name": model_name_pretty, "model_availability": model_availability, "urls": urls, "context_size": context_size, "submitted_by": submitted_by, } def _upload_predictions( self, task_id: str, model_folder: str, filenames: List[str], ) -> List[CommitOperationAdd]: commit_operations = [ CommitOperationAdd( path_in_repo=f"{task_id}/predictions/{model_folder}/{os.path.basename(filename)}", path_or_fileobj=filename, ) for filename in filenames ] return commit_operations def _compute_metrics_for_predictions( self, task_id: str, filenames: Optional[List[str]], temp_directory: str ) -> None: metrics_module = METRICS[task_id] assert ( metrics_module is not None ), f"Computing metrics for {task_id} is not supported." metrics_module.reset() open(os.path.join(temp_directory, "metrics.jsonl"), "w").close() # compute the metrics for each submitted file for filename in filenames: with jsonlines.open(filename, "r") as reader: for example in tqdm( reader, desc=f"Computing metrics for {os.path.basename(filename)}" ): metrics_module.add_batch( predictions=[example["prediction"]], references=[example["reference"]], ) computed_metrics = metrics_module.compute() metrics_module.reset() with jsonlines.open( os.path.join(temp_directory, "metrics.jsonl"), "a" ) as writer: writer.write(computed_metrics) # aggregate the metrics over submitted files with jsonlines.open( os.path.join(temp_directory, "metrics.jsonl"), "r" ) as reader: metrics_results = [line for line in reader] final_metrics_results = { key: sum(entry[key] for entry in metrics_results) / len(metrics_results) for key in metrics_results[0] } with open(os.path.join(temp_directory, "final_metrics.json"), "w") as f: json.dump(final_metrics_results, f) def _upload_results( self, task_id: str, model_folder: str, model_name_pretty: str, model_availability: str, urls: str, context_size: str, submitted_by: str, temp_directory: str, ) -> List[CommitOperationAdd]: 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(metadata_dict) with jsonlines.open( os.path.join(temp_directory, "final_results.jsonl"), "w" ) as writer: writer.write(final_results) return [ CommitOperationAdd( path_in_repo=f"{task_id}/results/{model_folder}.jsonl", path_or_fileobj=os.path.join(temp_directory, "final_results.jsonl"), ) ] def _verify_arguments( self, model_folder: str, model_name_pretty: str, model_availability: str, urls: str, context_size: str, submitted_by: str, filenames: Optional[List[str]], ): 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." assert ( model_availability ), "Please, specify non-empty information about a model's availability." assert ( context_size ), "Please, specify non-empty information about a model's context size." try: _ = int(context_size) except: raise ValueError( "Please, specify a model's context size as an integer (e.g., 16000)." ) assert ( submitted_by ), "Please, specify non-empty information about a submission's author(s)." assert filenames, "Please, attach at least one file with predictions." def upload_files( self, task_pretty: str, model_folder: str, model_name_pretty: str, model_availability: str, urls: str, context_size: str, submitted_by: str, filenames: Optional[List[str]], force: bool = False, ) -> str: try: self._verify_arguments( model_folder=model_folder, model_name_pretty=model_name_pretty, model_availability=model_availability, urls=urls, context_size=context_size, submitted_by=submitted_by, filenames=filenames, ) pr_title = f"🚀 New submission to {task_pretty} task: {model_name_pretty} with {context_size} context size from {submitted_by}" task_id = TASKS_PRETTY_REVERSE[task_pretty] logging.info("Checking if this request is already submitted...") if not force: if model_name_pretty in self._fs.ls( f"datasets/{self._dataset_id}/{task_id}/predictions" ) and all( filename in self._fs.ls( f"datasets/{self._dataset_id}/{task_id}/predictions/{model_name_pretty}" ) for filename in filenames + ["metadata.json"] ): return styled_warning( f"{model_name_pretty} is already present in {self._dataset_id}." ) 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}." ) logging.info("Processing predictions...") predictions_commit_operations = self._upload_predictions( task_id=task_id, model_folder=model_folder, filenames=filenames, ) with TemporaryDirectory() as d: logging.info("Computing metrics...") self._compute_metrics_for_predictions( task_id=task_id, filenames=filenames, temp_directory=str(d) ) logging.info("Processing results...") results_commit_operations = self._upload_results( task_id=task_id, model_folder=model_folder, model_name_pretty=model_name_pretty, model_availability=model_availability, urls=urls, context_size=context_size, submitted_by=submitted_by, temp_directory=str(d), ) logging.info("Creating commit...") new_pr = self._api.create_commit( repo_id=self._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! * Model name: {model_name_pretty} * Model availability: {model_availability} * Context Size: {context_size} * Relevant URLs: {urls} * Submitted By: {submitted_by} """, create_pr=True, repo_type="dataset", ) return styled_message(f"🎉 PR created at {new_pr.pr_url}.") except Exception as e: logging.exception(e) if str(e): return styled_error(str(e)) return styled_error("An exception occured.")