add llm-aggrefact experiment notebook

truera · sfc-gh-dhuang · Dec 10, 2024 · Oct 29, 2024 · Oct 31, 2024 · Oct 31, 2024
commit 35432e79a0bfc293b5c1e50b37ffa95cb6903057
diff --git a/src/benchmark/trulens/benchmark/benchmark_frameworks/experiments/dataset_preprocessing.py b/src/benchmark/trulens/benchmark/benchmark_frameworks/experiments/dataset_preprocessing.py
@@ -5,6 +5,7 @@
 import random
 from typing import Any, List, Tuple
 
+from datasets import load_dataset
 import ir_datasets
 import matplotlib.pyplot as plt
 import pandas as pd
@@ -13,6 +14,43 @@
 from trulens.feedback import GroundTruthAggregator
 
 
+def generate_balanced_llm_aggrefact_benchmark(split="test", random_seed=42):
+    llm_aggrefact_dataset = load_dataset("lytang/LLM-AggreFact")
+
+    # Convert to pandas DataFrame
+    df = pd.DataFrame(llm_aggrefact_dataset[split])
+
+    # Initialize an empty list to store balanced DataFrames
+    balanced_dfs = []
+
+    # Iterate over each unique dataset
+    for dataset_name in df["dataset"].unique():
+        # Filter the DataFrame for the current dataset
+        df_subset = df[df["dataset"] == dataset_name]
+
+        # Count the number of instances for each class
+        class_counts = df_subset["label"].value_counts()
+
+        # Determine the minimum count between the two classes
+        min_count = class_counts.min()
+
+        # Sample min_count instances from each class
+        df_balanced = (
+            df_subset.groupby("label")
+            .apply(lambda x: x.sample(min_count, random_state=random_seed))
+            .reset_index(drop=True)
+        )
+
+        # Append the balanced DataFrame to the list
+        balanced_dfs.append(df_balanced)
+
+    # Concatenate all balanced DataFrames into a final DataFrame
+    final_balanced_df = pd.concat(balanced_dfs, ignore_index=True)
+
+    # Display the balanced DataFrame
+    return final_balanced_df
+
+
 def generate_summeval_groundedness_golden_set(
     file_path: str, max_samples_per_bucket: int = 200
 ):

diff --git a/.../benchmark_frameworks/experiments/groundedness/llm_aggrefact_groundedness_benchmark.ipynb b/.../benchmark_frameworks/experiments/groundedness/llm_aggrefact_groundedness_benchmark.ipynb
diff --git a/...ens/benchmark/benchmark_frameworks/experiments/llm_aggrefact_groundedness_benchmark.ipynb b/...ens/benchmark/benchmark_frameworks/experiments/llm_aggrefact_groundedness_benchmark.ipynb
diff --git a/src/core/trulens/core/feedback/feedback.py b/src/core/trulens/core/feedback/feedback.py
@@ -946,7 +946,7 @@ def run(
 
                 else:
                     assert isinstance(
-                        result_val, (int, float, list)
+                        result_val, (int, float, list, dict)
                     ), f"Feedback function output must be a float or an int, a list of floats, or dict but was {type(result_val)}."
                     feedback_call = feedback_schema.FeedbackCall(
                         args=ins, ret=result_val, meta=meta