Add metrics computation for CMG task

app.py

@@ -1,6 +1,6 @@
 import os
 
-import gradio as gr
+import gradio as gr  # type: ignore[import]
 
 from src.content import (INTRODUCTION_TEXT, INTRODUCTION_TITLE,
                          LEADERBOARD_TEXT, LEADERBOARD_TITLE,

requirements.txt

@@ -1 +1,5 @@
-huggingface_hub
+huggingface_hub
+jsonlines
+pandas
+tqdm
+evaluate

src/evaluation/__init__.py

@@ -0,0 +1,3 @@
+from .metrics import METRICS
+
+__all__ = ["METRICS"]

src/evaluation/base_task_metrics.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List
+
+
+class BaseTaskMetrics(ABC):
+    def reset(self):
+        pass
+
+    @abstractmethod
+    def add_batch(
+        self, predictions: List[str], references: List[str], *args, **kwargs
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def compute(self, *args, **kwargs) -> Dict[str, float]:
+        pass

src/evaluation/commit_message_generation/__init__.py

@@ -0,0 +1,3 @@
+from .cmg_metrics import CMGMetrics
+
+__all__ = ["CMGMetrics"]

src/evaluation/commit_message_generation/cmg_metrics.py

@@ -0,0 +1,53 @@
+from typing import Dict, List
+
+import evaluate  # type: ignore[import]
+
+from ..base_task_metrics import BaseTaskMetrics
+from .b_norm import BNorm
+
+
+class CMGMetrics(BaseTaskMetrics):
+    def __init__(self):
+        self.bnorm = BNorm()
+        self.bleu = evaluate.load("sacrebleu")
+        self.chrf = evaluate.load("chrf")
+        self.rouge = evaluate.load("rouge")
+        self.bertscore = evaluate.load("bertscore")
+        self.bertscore_normalized = evaluate.load("bertscore")
+
+    def reset(self):
+        self.bnorm.reset()
+
+    def update(
+        self, predictions: List[str], references: List[str], *args, **kwargs
+    ) -> None:
+        self.bnorm.update(predictions=predictions, references=references)
+        self.bleu.add_batch(
+            predictions=predictions, references=[[ref] for ref in references]
+        )
+        self.chrf.add_batch(
+            predictions=predictions, references=[[ref] for ref in references]
+        )
+        self.rouge.add_batch(predictions=predictions, references=references)
+        self.bertscore.add_batch(predictions=predictions, references=references)
+        self.bertscore_normalized.add_batch(
+            predictions=predictions, references=references
+        )
+
+    def compute(self, *args, **kwargs) -> Dict[str, float]:
+        rouge = self.rouge.compute()
+        bertscore = self.bertscore.compute(lang="en")
+        bertscore_normalized = self.bertscore_normalized.compute(
+            lang="en", rescale_with_baseline=True
+        )
+        return {
+            "bnorm": self.bnorm.compute(),
+            "bleu": self.bleu.compute(tokenize="13a")["score"],
+            "chrf": self.chrf.compute()["score"],
+            "rouge1": rouge["rouge1"] * 100,
+            "rouge2": rouge["rouge2"] * 100,
+            "rougeL": rouge["rougeL"] * 100,
+            "bertscore": sum(bertscore["f1"]) / len(bertscore["f1"]),
+            "bertscore_normalized": sum(bertscore_normalized["f1"])
+            / len(bertscore_normalized["f1"]),
+        }

src/evaluation/metrics.py

@@ -0,0 +1,13 @@
+from typing import Dict, Optional
+
+from .base_task_metrics import BaseTaskMetrics
+from .commit_message_generation import CMGMetrics
+
+METRICS: Dict[str, Optional[BaseTaskMetrics]] = {
+    "commit_message_generation": CMGMetrics(),
+    "bug_localization": None,
+    "module_to_text": None,
+    "library_usage": None,
+    "project_code_completion": None,
+    "bug_localization_build_logs": None,
+}

src/formatting.py

@@ -0,0 +1,12 @@
+def styled_error(error):
+    return f"<p style='color: red; font-size: 20px; text-align: center;'>{error}</p>"
+
+
+def styled_warning(warn):
+    return f"<p style='color: orange; font-size: 20px; text-align: center;'>{warn}</p>"
+
+
+def styled_message(message):
+    return (
+        f"<p style='color: green; font-size: 20px; text-align: center;'>{message}</p>"
+    )

src/get_results_for_task.py

@@ -1,6 +1,4 @@
-import pandas as pd
-
-RESULTS_DATASET = "JetBrains-Research/lca-results"
+import pandas as pd  # type: ignore[import]
 
 
 def get_results_for_task_stub(task: str) -> pd.DataFrame:

src/submission_uploader.py

@@ -1,9 +1,16 @@
 import json
+import logging
 import os
-from typing import List, Optional
+from tempfile import TemporaryDirectory
+from typing import Dict, List, Optional
 
-from huggingface_hub import CommitOperationAdd, Discussion, HfApi, HfFileSystem
+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
 
 
@@ -39,19 +46,30 @@ class SubmissionUploader:
                 and discussion.title == pr_title
             ):
                 return discussion
+        return None
 
-    def _upload_files(
+    def _get_metadata(
         self,
-        task_id: str,
-        model_folder: str,
         model_name_pretty: str,
         model_availability: str,
         urls: str,
         context_size: str,
         submitted_by: str,
-        filenames: Optional[List[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]:
-        # add predictions files
         commit_operations = [
             CommitOperationAdd(
                 path_in_repo=f"{task_id}/predictions/{model_folder}/{os.path.basename(filename)}",
@@ -59,25 +77,114 @@ class SubmissionUploader:
             )
             for filename in filenames
         ]
+        return commit_operations
 
-        # add metadata file
-        metadata_dict = {
-            "model_name": model_name_pretty,
-            "model_availability": model_availability,
-            "urls": urls,
-            "context_size": context_size,
-            "submitted_by": submitted_by,
+    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("metadata.json", "w") as f:
-            json.dump(metadata_dict, f)
-        commit_operations.append(
+        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}/predictions/{model_folder}/metadata.json",
-                path_or_fileobj="metadata.json",
+                path_in_repo=f"{task_id}/results/{model_folder}.jsonl",
+                path_or_fileobj=os.path.join(temp_directory, "final_results.jsonl"),
             )
-        )
+        ]
 
-        return commit_operations
+    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,
@@ -92,10 +199,21 @@ class SubmissionUploader:
         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"
@@ -106,29 +224,46 @@ class SubmissionUploader:
                     )
                     for filename in filenames + ["metadata.json"]
                 ):
-                    return (
+                    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 f"{self._dataset_id} already has an open PR for this submission: {url}."
+                    return styled_warning(
+                        f"{self._dataset_id} already has an open PR for this submission: {url}."
+                    )
 
-            commit_operations = self._upload_files(
+            logging.info("Processing predictions...")
+            predictions_commit_operations = self._upload_predictions(
                 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,
                 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=commit_operations,
+                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!
                     
@@ -141,7 +276,10 @@ class SubmissionUploader:
                 create_pr=True,
                 repo_type="dataset",
             )
-            return f"🎉 PR created at {new_pr.pr_url}."
+            return styled_message(f"🎉 PR created at {new_pr.pr_url}.")
 
-        except Exception:
-            return "An exception occured."
+        except Exception as e:
+            logging.exception(e)
+            if str(e):
+                return styled_error(str(e))
+            return styled_error("An exception occured.")