edit test.py with MetricInput

Marker-Inc-Korea · vkehfdl1 · Sep 11, 2024 · Aug 26, 2024 · Sep 1, 2024 · Sep 1, 2024
commit 513a56a31d95ec4ee51c6349dc06d006edb8d257
diff --git a/tests/autorag/evaluate/metric/test_generation_metric.py b/tests/autorag/evaluate/metric/test_generation_metric.py
@@ -4,6 +4,7 @@
 from llama_index.embeddings.openai import OpenAIEmbedding
 
 from autorag.evaluation.metric import bleu, meteor, rouge, sem_score, g_eval, bert_score
+from autorag.schema.metricinput import MetricInput
 from tests.delete_tests import is_github_action
 from tests.mock import mock_get_text_embedding_batch
 
@@ -38,9 +39,12 @@
 	"요즘 세상에서는 예술가가 되려면, AI를 이겨야 한다.",
 ]
 
-
+metric_inputs = [MetricInput(generated_texts=gen, generation_gt=gen_gt) for gen, gen_gt in
+				 zip(generations, generation_gts)]
+ko_metric_inputs = [MetricInput(generated_texts=gen, generation_gt=gen_gt) for gen, gen_gt in
+					zip(ko_generations, ko_generation_gts)]
 def base_test_generation_metrics(func, solution, **kwargs):
-	scores = func(generation_gt=generation_gts, generations=generations, **kwargs)
+	scores = func(metric_inputs, **kwargs)
 	assert len(scores) == len(generation_gts)
 	assert all(isinstance(score, float) for score in scores)
 	assert all(
@@ -49,7 +53,7 @@ def base_test_generation_metrics(func, solution, **kwargs):
 
 
 def ko_base_test_generation_metrics(func, solution, **kwargs):
-	scores = func(generation_gt=ko_generation_gts, generations=ko_generations, **kwargs)
+	scores = func(ko_metric_inputs, **kwargs)
 	assert len(scores) == len(ko_generation_gts)
 	assert all(isinstance(score, float) for score in scores)
 	assert all(
@@ -86,8 +90,7 @@ def test_sem_score():
 )
 def test_sem_score_other_model():
 	scores = sem_score(
-		generation_gt=generation_gts,
-		generations=generations,
+		metric_inputs=metric_inputs,
 		embedding_model=OpenAIEmbedding(),
 	)
 	assert len(scores) == len(generation_gts)

diff --git a/tests/autorag/evaluate/metric/test_retrieval_contents_metric.py b/tests/autorag/evaluate/metric/test_retrieval_contents_metric.py
@@ -6,6 +6,7 @@
 	retrieval_token_precision,
 	retrieval_token_recall,
 )
+from autorag.schema.metricinput import MetricInput
 
 gt = [
 	["Enough for drinking water", "Just looking for a water bottle"],
@@ -23,6 +24,7 @@
 	["Who is son? He is great player in the world"],
 	["i love havertz", "i love kai havertz"],
 ]
+metric_inputs = [MetricInput(gt_contents=g, retrieval_contents=p) for g, p in zip(gt, pred)]
 
 
 def test_single_token_f1():
@@ -38,23 +40,23 @@ def test_single_token_f1():
 
 
 def test_retrieval_token_f1():
-	f1 = retrieval_token_f1.__wrapped__(gt[0], pred[0])
+	f1 = retrieval_token_f1.__wrapped__(MetricInput(gt_contents=gt[0], retrieval_contents=pred[0]))
 	assert f1 == pytest.approx(0.38333, rel=0.001)
 
-	f1 = retrieval_token_f1.__wrapped__(gt[1], pred[1])
+	f1 = retrieval_token_f1.__wrapped__(MetricInput(gt_contents=gt[1], retrieval_contents=pred[1]))
 	assert f1 == pytest.approx(0.797979, rel=0.001)
 
-	result_f1 = retrieval_token_f1(gt_contents=gt, pred_contents=pred)
+	result_f1 = retrieval_token_f1(metric_inputs=metric_inputs)
 	assert result_f1 == pytest.approx([0.38333, 0.797979, None, None], rel=0.001)
 
 
 def test_retrieval_token_precision():
-	result_precision = retrieval_token_precision(gt_contents=gt, pred_contents=pred)
+	result_precision = retrieval_token_precision(metric_inputs=metric_inputs)
 	assert result_precision == pytest.approx(
 		[0.383333, 0.8222222, None, None], rel=0.001
 	)
 
 
 def test_retrieval_token_recall():
-	result_recall = retrieval_token_recall(gt_contents=gt, pred_contents=pred)
+	result_recall = retrieval_token_recall(metric_inputs=metric_inputs)
 	assert result_recall == pytest.approx([0.383333, 0.777777, None, None], rel=0.001)
diff --git a/tests/autorag/evaluate/metric/test_retrieval_metric.py b/tests/autorag/evaluate/metric/test_retrieval_metric.py
@@ -9,6 +9,7 @@
 	retrieval_mrr,
 	retrieval_map,
 )
+from autorag.schema.metricinput import MetricInput
 
 retrieval_gt = [
 	[["test-1", "test-2"], ["test-3"]],
@@ -36,11 +37,11 @@
 	["pred-14"],  # retrieval_gt is empty so not counted
 	["pred-15", "pred-16", "test-15"],  # recall:1, precision: 1/3, f1: 0.5
 ]
-
+metric_inputs = [MetricInput(retrieval_gt=ret_gt, retrieval_ids=pr) for ret_gt, pr in zip(retrieval_gt, pred)]
 
 def test_retrieval_f1():
 	solution = [0.5, 2 / 7, 2 / 5, 4 / 7, 2 / 3, None, None, 0.5]
-	result = retrieval_f1(retrieval_gt=retrieval_gt, pred_ids=pred)
+	result = retrieval_f1(metric_inputs=metric_inputs)
 	for gt, res in zip(solution, result):
 		assert gt == pytest.approx(res, rel=1e-4)
 
@@ -49,21 +50,23 @@ def test_numpy_retrieval_metric():
 	retrieval_gt_np = [[np.array(["test-1", "test-4"])], np.array([["test-2"]])]
 	pred_np = np.array([["test-2", "test-3", "test-1"], ["test-5", "test-6", "test-8"]])
 	solution = [1.0, 0.0]
-	result = retrieval_recall(retrieval_gt=retrieval_gt_np, pred_ids=pred_np)
+	metric_inputs_np = [MetricInput(retrieval_gt=ret_gt_np, retrieval_ids=pr_np) for ret_gt_np, pr_np in
+						zip(retrieval_gt_np, pred_np)]
+	result = retrieval_recall(metric_inputs=metric_inputs_np)
 	for gt, res in zip(solution, result):
 		assert gt == pytest.approx(res, rel=1e-4)
 
 
 def test_retrieval_recall():
 	solution = [0.5, 1 / 3, 1, 2 / 3, 1, None, None, 1]
-	result = retrieval_recall(retrieval_gt=retrieval_gt, pred_ids=pred)
+	result = retrieval_recall(metric_inputs=metric_inputs)
 	for gt, res in zip(solution, result):
 		assert gt == pytest.approx(res, rel=1e-4)
 
 
 def test_retrieval_precision():
 	solution = [0.5, 0.25, 0.25, 0.5, 0.5, None, None, 1 / 3]
-	result = retrieval_precision(retrieval_gt=retrieval_gt, pred_ids=pred)
+	result = retrieval_precision(metric_inputs=metric_inputs)
 	for gt, res in zip(solution, result):
 		assert gt == pytest.approx(res, rel=1e-4)
 
@@ -79,20 +82,20 @@ def test_retrieval_ndcg():
 		None,
 		0.5,
 	]
-	result = retrieval_ndcg(retrieval_gt=retrieval_gt, pred_ids=pred)
+	result = retrieval_ndcg(metric_inputs=metric_inputs)
 	for gt, res in zip(solution, result):
 		assert gt == pytest.approx(res, rel=1e-4)
 
 
 def test_retrieval_mrr():
 	solution = [1 / 2, 1 / 3, 1, 1 / 2, 1, None, None, 1 / 3]
-	result = retrieval_mrr(retrieval_gt=retrieval_gt, pred_ids=pred)
+	result = retrieval_mrr(metric_inputs=metric_inputs)
 	for gt, res in zip(solution, result):
 		assert gt == pytest.approx(res, rel=1e-4)
 
 
 def test_retrieval_map():
 	solution = [5 / 12, 1 / 3, 1, 1 / 2, 1, None, None, 1 / 3]
-	result = retrieval_map(retrieval_gt=retrieval_gt, pred_ids=pred)
+	result = retrieval_map(metric_inputs=metric_inputs)
 	for gt, res in zip(solution, result):
 		assert gt == pytest.approx(res, rel=1e-4)