data analysis notebooks

llm_toolkit/translation_utils.py

@@ -163,6 +163,24 @@ def load_translation_dataset(data_path, tokenizer=None):
     return datasets
 
 
+def count_entries_with_max_tokens(entries, max_tokens):
+    """
+    Count the number of entries with the max output tokens or more.
+
+    Parameters:
+    entries (list of int): List of token counts for each entry.
+    max_tokens (int): The maximum token threshold.
+
+    Returns:
+    int: The number of entries with token counts greater than or equal to max_tokens.
+    """
+    count = 0
+    for tokens in entries:
+        if tokens >= max_tokens:
+            count += 1
+    return count
+
+
 def get_metrics(df, max_output_tokens=2048):
     metrics_df = pd.DataFrame(df.columns.T)[2:]
     metrics_df.rename(columns={0: "model"}, inplace=True)
@@ -203,15 +221,15 @@ def get_metrics(df, max_output_tokens=2048):
         )
 
         num_entries_with_max_output_tokens.append(
-            df["output_tokens"].value_counts().get(max_output_tokens, 0)
+            count_entries_with_max_tokens(df["output_tokens"], max_output_tokens)
         )
 
     metrics_df["meteor"] = meteor
     metrics_df["bleu_1"] = bleu_1
     metrics_df["rouge_l"] = rouge_l
     metrics_df["ews_score"] = ews_score
-    metrics_df["repetition_score"] = ews_score
-    metrics_df["total_repetitions"] = ews_score
+    metrics_df["repetition_score"] = repetition_score
+    metrics_df["total_repetitions"] = total_repetitions
     metrics_df["num_entries_with_max_output_tokens"] = (
         num_entries_with_max_output_tokens
     )