divide results

src/evaluate/evaluate_answers.py

@@ -66,22 +66,24 @@ async def deepseek_eval(system: str, prompt: str, eval_model: str, host: str) ->
         return completion
 
 
-async def generate(df: List, cand: str, base: str, eval_model: str, host: str) -> tuple:
+async def generate(df_cand: str, df_base: str, eval_model: str, host: str) -> tuple:
     tasks = []
     models_answers = []
+    base, cand = df_cand[0]["replies"]["model_name"], df_base[0]["replies"]["model_name"]
 
-    for instance in df:
+    for instance_cand, instance_base in tqdm(zip(df_cand, df_base)):
         # extract only the responses from the input dict
         responses = {base: None, cand: None}
-        for resp in instance["replies"]:
-            if resp["model_name"] == base or resp["model_name"] == cand:
-                responses[resp["model_name"]] = resp["text"]
+        if instance_cand["replies"]["model_name"] == cand:
+            responses[instance_cand["replies"]["model_name"]] = instance_cand["replies"]["text"]
+        if instance_base["replies"]["model_name"] == base:
+            responses[instance_base["replies"]["model_name"]] = instance_base["replies"]["text"]
 
         if responses[base] is None or responses[cand] is None:
             raise ValueError("There are no cand or base model answer")
 
         prompt = "Question: {}\n\nFirst Response: {}\nSecond Response: {}\nAnswer: ".format(
-            instance["turns"]["content"],
+            instance_cand["turns"]["content"],
             responses[cand],
             responses[base],
         )
@@ -93,22 +95,24 @@ async def generate(df: List, cand: str, base: str, eval_model: str, host: str) -
     return models_answers, eval_results
 
 
-def generate_sync(df: List, cand: str, base: str, eval_model: str, host: str) -> tuple:
+def generate_sync(df_cand: str, df_base: str, eval_model: str, host: str) -> tuple:
     models_answers = []
     eval_results = []
+    base, cand = df_cand[0]["replies"]["model_name"], df_base[0]["replies"]["model_name"]
 
-    for instance in tqdm(df):
+    for instance_cand, instance_base in tqdm(zip(df_cand, df_base)):
         # extract only the responses from the input dict
         responses = {base: None, cand: None}
-        for resp in instance["replies"]:
-            if resp["model_name"] == base or resp["model_name"] == cand:
-                responses[resp["model_name"]] = resp["text"]
+        if instance_cand["replies"]["model_name"] == cand:
+            responses[instance_cand["replies"]["model_name"]] = instance_cand["replies"]["text"]
+        if instance_base["replies"]["model_name"] == base:
+            responses[instance_base["replies"]["model_name"]] = instance_base["replies"]["text"]
 
         if responses[base] is None or responses[cand] is None:
             raise ValueError("There are no cand or base model answer")
 
         prompt = "Question: {}\n\nFirst Response: {}\nSecond Response: {}\nAnswer: ".format(
-            instance["turns"]["content"],
+            instance_cand["turns"]["content"],
             responses[cand],
             responses[base],
         )
@@ -205,20 +209,22 @@ def draw_dataframe(
 
 def get_df(
     input_filename: str,
-    output_filename: str,
+    output_filename: str = None,
 ):
 
     with open(input_filename, "r") as f:
         df = [json.loads(line) for line in f]
 
-    if os.path.exists(output_filename):
-        os.remove(output_filename)
+    if output_filename:
+        if os.path.exists(output_filename):
+            os.remove(output_filename)
 
     return df
 
 
 def main(
-    input_filename: str,
+    data_root_dir: str,
+    dataset_name: str,
     output_filename: str,
     baseline_model: str,
     candidate_model: str,
@@ -228,26 +234,32 @@ def main(
     host: str,
     sync: bool = False,
 ):
-    df = get_df(
-        input_filename=input_filename,
+    dataset_name_cut = dataset_name.split("/")[-1]
+    input_cand_filename = os.path.join(data_root_dir, "generations", f"{candidate_model}_{dataset_name_cut}_responses.jsonl")
+    input_base_filename = os.path.join(data_root_dir, "generations", f"{baseline_model}_{dataset_name_cut}_responses.jsonl")
+    output_filename = os.path.join(data_root_dir, "judgements", f"{candidate_model}_vs_{baseline_model}_{dataset_name_cut}.jsonl")
+    df_cand = get_df(
+        input_filename=input_cand_filename,
         output_filename=output_filename,
     )
+    df_base = get_df(
+        input_filename=input_base_filename,
+    )
 
     if swap:
-        baseline_model, candidate_model = candidate_model, baseline_model
+        df_cand, df_base = df_base, df_cand
     if sync:
         models_answers, eval_results = generate_sync(
-            df=df, cand=candidate_model, base=baseline_model, eval_model=eval_model, host=host
+            df_cand=df_cand, df_base=df_base, eval_model=eval_model, host=host
         )
     else:
         models_answers, eval_results = [], []
-        bar = tqdm(total=len(df))
-        for i in range(0, len(df), chunk_size):
+        bar = tqdm(total=len(df_cand))
+        for i in range(0, len(df_cand), chunk_size):
             model_answer, eval_result = asyncio.run(
                 generate(
-                    df=df[i:i + chunk_size],
-                    cand=candidate_model,
-                    base=baseline_model,
+                    df_cand=df_cand[i:i + chunk_size],
+                    df_base=df_base[i:i + chunk_size],
                     eval_model=eval_model,
                     host=host,
                 )
@@ -267,7 +279,7 @@ def main(
     cat_only_accuracy = defaultdict(int)
     cat_tie_accuracy = defaultdict(int)
 
-    for instance, eval_result, answers in zip(df, eval_results, models_answers):
+    for instance, eval_result, answers in zip(df_cand, eval_results, models_answers):
         if eval_result is None:
             total -= 1
             continue
@@ -348,7 +360,7 @@ def main(
         candidate_model=candidate_model,
         four_numbers=[accuracy, other_accuracy, both_good_accuracy, both_bad_accuracy],
         two_numbers=pao_to_2,
-        input_filename=input_filename,
+        input_filename=input_cand_filename,
         swap=swap,
     )
 
@@ -364,10 +376,15 @@ if __name__ == "__main__":
         help="Хостнейм, на котором крутится модель",
     )
     parser.add_argument(
-        "--input-filename",
+        "--data-root-dir",
         type=str,
         help="Файл который надо оценить",
     )
+    parser.add_argument(
+        "--dataset-name",
+        type=str,
+        help="Название бенчмарка",
+    )
     parser.add_argument("--output-filename", type=str, default="judge_results_with_probs.jsonl")
     parser.add_argument(
         "--sleep-time",
@@ -416,7 +433,7 @@ if __name__ == "__main__":
 
     for i in range(len(swaps)):
         results[swaps[i]] = main(
-            input_filename=args.input_filename,
+            data_root_dir=args.data_root_dir,
             output_filename=f"swap_{i}_{args.candidate_model}",
             baseline_model=args.baseline_model,
             candidate_model=args.candidate_model,

src/evaluate/generate_answers.py

@@ -30,16 +30,13 @@ def write_response_jsonl(response_text, counter, question, model_name, output_fi
         "question_id": question["question_id"][counter],
         "cluster": question["cluster"][counter],
         "turns": question["turns"][counter],
-        "replies": question.get("replies", [])
         }
 
-    cur_dict["replies"].append(
-        {
+    cur_dict["replies"] = {
             "message_id": message_id,
             "text": response_text,
             "model_name": model_name,
-        }
-    )
+            }
 
     with open(output_filename, "a") as f:
         json.dump(cur_dict, f, ensure_ascii=False)