style: code blacked

app.py

@@ -52,7 +52,9 @@ with gr.Blocks(
     ),
 ) as application:
     gr.Markdown("# 🥇 ROMB - Russian Olympiad Math Benchmark")
-    gr.Markdown(f"See ROMB-1.0 dataset there - [{DATASET_NAME}](https://huggingface.co/datasets/{DATASET_NAME}).")
+    gr.Markdown(
+        f"See ROMB-1.0 dataset there - [{DATASET_NAME}](https://huggingface.co/datasets/{DATASET_NAME})."
+    )
     with gr.Tabs():
         with gr.Tab("Leaderboard"):
             gr.Markdown("In progress...")
@@ -77,7 +79,7 @@ with gr.Blocks(
                     [22],
                     [40],
                     [230],
-                ]
+                ],
             )
 
 

src/common/data.py

@@ -13,5 +13,7 @@ def load_dataset() -> pd.DataFrame:
     ds = datasets.load_dataset(DATASET_NAME, split="test")
     df = pd.DataFrame(ds)
 
-    df[DatasetSchema.correct_answer] = df[DatasetSchema.correct_answer].apply(json.loads)
+    df[DatasetSchema.correct_answer] = df[DatasetSchema.correct_answer].apply(
+        json.loads
+    )
     return df

src/eval/cli.py

@@ -41,7 +41,9 @@ def _evaluate_single_answer(
         )
     except Exception as e:
         print(e)
-        print(f"Error evaluating row with {row[DatasetSchema.check_type]} {row[DatasetSchema.id_]}: {y_true} vs {y_pred}")
+        print(
+            f"Error evaluating row with {row[DatasetSchema.check_type]} {row[DatasetSchema.id_]}: {y_true} vs {y_pred}"
+        )
         exit(1)
     return result
 
@@ -53,7 +55,9 @@ def _evaluate(
 ) -> pd.DataFrame:
     tqdm.pandas()
 
-    generated_df[GeneratedDatasetSchema.generated_answer] = generated_df[GeneratedDatasetSchema.generated_answer].apply(
+    generated_df[GeneratedDatasetSchema.generated_answer] = generated_df[
+        GeneratedDatasetSchema.generated_answer
+    ].apply(
         lambda x: GenerationAnswer.model_validate(deepcopy(x)) if x else None,
     )
     dataset_df = load_dataset()
@@ -67,10 +71,14 @@ def _evaluate(
         axis=1,
     )
 
-    predictions_df[DatasetEvalSchema.predicted_answer] = predictions_df[GeneratedDatasetSchema.generated_answer].apply(
+    predictions_df[DatasetEvalSchema.predicted_answer] = predictions_df[
+        GeneratedDatasetSchema.generated_answer
+    ].apply(
         lambda x: x.answer if not pd.isna(x) else None,
     )
-    predictions_df[DatasetEvalSchema.context] = predictions_df[GeneratedDatasetSchema.generated_answer].apply(
+    predictions_df[DatasetEvalSchema.context] = predictions_df[
+        GeneratedDatasetSchema.generated_answer
+    ].apply(
         lambda x: x.context if not pd.isna(x) else None,
     )
     predictions_df = predictions_df[list(DatasetEvalSchema._collect_fields().keys())]
@@ -91,7 +99,9 @@ def evaluate(
 
     df = pd.read_json(file, lines=True)
     evaluated_df = _evaluate(df)
-    evaluated_df.to_json(file.with_suffix(".eval.jsonl"), orient="records", lines=True, force_ascii=False)
+    evaluated_df.to_json(
+        file.with_suffix(".eval.jsonl"), orient="records", lines=True, force_ascii=False
+    )
 
 
 @pa.check_input(DatasetEvalSchema)
@@ -101,16 +111,24 @@ def _metrics(
     model_name: str,
     model_size: float,
     model_url: str,
-    model_config: str
+    model_config: str,
 ) -> pd.DataFrame:
     pass1 = df[DatasetEvalSchema.is_correct].mean()
 
     w = df[DatasetEvalSchema.grade].apply(grade_to_weight)
-    weighted_accuracy = (df[DatasetEvalSchema.is_correct].astype(int) * w).sum() / w.sum()
-
-    arith_pass1 = df[df[DatasetEvalSchema.task_type] == "arith"][DatasetEvalSchema.is_correct].mean()
-    geometry_pass1 = df[df[DatasetEvalSchema.task_type] == "geometry"][DatasetEvalSchema.is_correct].mean()
-    logic_pass1 = df[df[DatasetEvalSchema.task_type] == "logic"][DatasetEvalSchema.is_correct].mean()
+    weighted_accuracy = (
+        df[DatasetEvalSchema.is_correct].astype(int) * w
+    ).sum() / w.sum()
+
+    arith_pass1 = df[df[DatasetEvalSchema.task_type] == "arith"][
+        DatasetEvalSchema.is_correct
+    ].mean()
+    geometry_pass1 = df[df[DatasetEvalSchema.task_type] == "geometry"][
+        DatasetEvalSchema.is_correct
+    ].mean()
+    logic_pass1 = df[df[DatasetEvalSchema.task_type] == "logic"][
+        DatasetEvalSchema.is_correct
+    ].mean()
 
     result = {
         LeaderBoardSchema.model_name: model_name,
@@ -174,7 +192,7 @@ def metrics(
         model_name=model_name,
         model_size=model_size,
         model_url=model_url,
-        model_config=model_config or '',
+        model_config=model_config or "",
     )
     metrics = metrics_df.to_dict(orient="records")[0]
     print(f"Metrics for {model_name}:")

src/eval/matchers.py

@@ -32,7 +32,9 @@ def um(y_true: list, y_pred: list) -> bool:
         return False
     if len(y_true) == 0:
         return True
-    if (len(y_true) > 0 and type(y_true[0]) is dict) or (len(y_true) == 0 and type(y_pred[0]) is dict):
+    if (len(y_true) > 0 and type(y_true[0]) is dict) or (
+        len(y_true) == 0 and type(y_pred[0]) is dict
+    ):
         y_true = [_dict_to_tuple(item) for item in y_true]
         y_pred = [_dict_to_tuple(item) for item in y_pred]
     if type(y_true) != type(y_pred):

src/eval/metrics.py

@@ -3,5 +3,5 @@ import numpy as np
 
 def grade_to_weight(g: str) -> float:
     """Convert a grade string to a weight value."""
-    parts = list(map(int, g.split('-')))
+    parts = list(map(int, g.split("-")))
     return np.mean(parts)

src/generate/cli.py

@@ -41,7 +41,9 @@ def _generate_single_answer(
 ) -> GenerationAnswer:
     if temp_path and (temp_path / f"{row[DatasetSchema.id_]}.json").exists():
         return GenerationAnswer.model_validate(
-            json.load(open(temp_path / f"{row[DatasetSchema.id_]}.json", "r"))[GeneratedDatasetSchema.generated_answer]
+            json.load(open(temp_path / f"{row[DatasetSchema.id_]}.json", "r"))[
+                GeneratedDatasetSchema.generated_answer
+            ]
         )
     answer_type = make_root_model(row[DatasetSchema.answer_type])
     chain = build_chain(answer_type)
@@ -131,14 +133,16 @@ def generate(
         build_chain_function,
         llm_class=config.llm_class,
         structured_output_method=config.structured_output_method,
-        **config.kwargs
+        **config.kwargs,
     )
 
-    df = _generate_answers(df, build_chain_function, use_tqdm=use_tqdm, temp_path=temp_path)
-
-    df[GeneratedDatasetSchema.generated_answer] = df[GeneratedDatasetSchema.generated_answer].apply(
-        lambda x: x.model_dump()
+    df = _generate_answers(
+        df, build_chain_function, use_tqdm=use_tqdm, temp_path=temp_path
     )
+
+    df[GeneratedDatasetSchema.generated_answer] = df[
+        GeneratedDatasetSchema.generated_answer
+    ].apply(lambda x: x.model_dump())
     df.to_json(
         output_path,
         lines=True,
@@ -151,15 +155,19 @@ def generate(
 def _type_sanitycheck(
     generated_df: pd.DataFrame,
 ) -> tuple[bool, str]:
-    generated_df[GeneratedDatasetSchema.generated_answer] = generated_df[GeneratedDatasetSchema.generated_answer].apply(
-        lambda x: GenerationAnswer.model_validate(deepcopy(x)) if not isinstance(x, GenerationAnswer) else x
+    generated_df[GeneratedDatasetSchema.generated_answer] = generated_df[
+        GeneratedDatasetSchema.generated_answer
+    ].apply(
+        lambda x: GenerationAnswer.model_validate(deepcopy(x))
+        if not isinstance(x, GenerationAnswer)
+        else x
     )
 
     dataset_df = load_dataset()
     predicted_df = dataset_df.join(
         generated_df.set_index(GeneratedDatasetSchema.id_),
         on=DatasetSchema.id_,
-        rsuffix='_generated',
+        rsuffix="_generated",
     ).dropna(subset=[GeneratedDatasetSchema.generated_answer])
 
     if len(predicted_df) == 0:
@@ -170,13 +178,20 @@ def _type_sanitycheck(
         lambda row: matches_type(
             row[GeneratedDatasetSchema.generated_answer].answer,
             string_to_type(row[DatasetSchema.answer_type]),
-        ), axis=1
+        ),
+        axis=1,
     )
 
     if not predicted_df[TYPE_MATCH].all():
-        return False, f"Type mismatch found for {predicted_df[~predicted_df[TYPE_MATCH]][DatasetSchema.id_].tolist()}."
+        return (
+            False,
+            f"Type mismatch found for {predicted_df[~predicted_df[TYPE_MATCH]][DatasetSchema.id_].tolist()}.",
+        )
 
-    return True, f"All matched. Predicted count: {len(predicted_df)} of {len(dataset_df)}"
+    return (
+        True,
+        f"All matched. Predicted count: {len(predicted_df)} of {len(dataset_df)}",
+    )
 
 
 @click.command()

src/generate/generators.py

@@ -57,9 +57,7 @@ def build_singleturn_chain(
             context={},
         )
     )
-    chain = chain.with_retry(
-        retry_if_exception_type=(openai.PermissionDeniedError, )
-    )
+    chain = chain.with_retry(retry_if_exception_type=(openai.PermissionDeniedError,))
     return chain
 
 
@@ -126,9 +124,7 @@ def build_thinking_chain(
             )
         )
     )
-    chain = chain.with_retry(
-        retry_if_exception_type=(openai.PermissionDeniedError, )
-    )
+    chain = chain.with_retry(retry_if_exception_type=(openai.PermissionDeniedError,))
     return chain