From ab0aa93fbcf474c38a062c47b414db1cd5eaeedb Mon Sep 17 00:00:00 2001
From: Kush Dubey <kushdubey63@gmail.com>
Date: Fri, 1 Nov 2024 00:37:10 -0700
Subject: [PATCH] Flat parametrize test args

---
 CITATION.cff                       |  6 +--
 LICENSE                            |  2 +-
 docs/source/conf.py                |  4 +-
 tests/utils/test_utils_classify.py | 73 ++++++++++++------------------
 4 files changed, 36 insertions(+), 49 deletions(-)

diff --git a/CITATION.cff b/CITATION.cff
index a730a51..dfe0a5f 100644
--- a/CITATION.cff
+++ b/CITATION.cff
@@ -17,6 +17,6 @@ abstract: >-
   Completion After Prompt Probability. Make your LLM make a
   choice
 license: Apache-2.0
-commit: a46dd22
-version: 0.9.0
-date-released: '2024-02-01'
+commit: ' 3abd478'
+version: 0.9.5
+date-released: '2024-10-28'
diff --git a/LICENSE b/LICENSE
index 49ba70c..8de3217 100644
--- a/LICENSE
+++ b/LICENSE
@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright 2023-2023 Kush Dubey
+   Copyright 2023-2024 Kush Dubey
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
diff --git a/docs/source/conf.py b/docs/source/conf.py
index e776637..5d9d434 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -14,8 +14,8 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
 project = "CAPPr"
-copyright = "2023, kddubey"
-author = "kddubey"
+copyright = "2023-2024, Kush Dubey"
+author = "Kush Dubey"
 
 version = __version__
 release = version
diff --git a/tests/utils/test_utils_classify.py b/tests/utils/test_utils_classify.py
index 727d54c..8fb28ee 100644
--- a/tests/utils/test_utils_classify.py
+++ b/tests/utils/test_utils_classify.py
@@ -87,31 +87,26 @@ def test_agg_log_probs():
         _ = classify.agg_log_probs([[[[0, 1]]]])
 
 
-@pytest.mark.parametrize("likelihoods", ([[4, 1], [1, 4]],))
-@pytest.mark.parametrize("prior", (None, [1 / 2, 1 / 2], [1 / 3, 2 / 3]))
-@pytest.mark.parametrize("normalize", (True, False))
-def test_posterior_prob_2d(likelihoods, prior, normalize):
-    # TODO: clean this up
+@pytest.mark.parametrize(
+    "likelihoods, prior, normalize, expected",
+    [
+        # Case 1: Prior = None
+        ([[4, 1], [1, 4]], None, True, [[4 / 5, 1 / 5], [1 / 5, 4 / 5]]),
+        ([[4, 1], [1, 4]], None, False, [[4, 1], [1, 4]]),
+        # Case 2: Prior = [1 / 2, 1 / 2]
+        ([[4, 1], [1, 4]], [1 / 2, 1 / 2], True, [[4 / 5, 1 / 5], [1 / 5, 4 / 5]]),
+        ([[4, 1], [1, 4]], [1 / 2, 1 / 2], False, [[2, 0.5], [0.5, 2]]),
+        # Case 3: Prior = [1 / 3, 2 / 3]
+        ([[4, 1], [1, 4]], [1 / 3, 2 / 3], True, [[2 / 3, 1 / 3], [1 / 9, 8 / 9]]),
+        ([[4, 1], [1, 4]], [1 / 3, 2 / 3], False, [[4 / 3, 2 / 3], [1 / 3, 8 / 3]]),
+    ],
+)
+def test_posterior_prob_2d(likelihoods, prior, normalize, expected):
     posteriors = classify.posterior_prob(
         likelihoods, axis=1, prior=prior, normalize=normalize
     )
-    if prior == [1 / 2, 1 / 2]:
-        if normalize:
-            assert np.all(np.isclose(posteriors, [[4 / 5, 1 / 5], [1 / 5, 4 / 5]]))
-        else:
-            assert np.all(posteriors == np.array(likelihoods) / 2)
-    elif prior is None:
-        if normalize:
-            assert np.all(np.isclose(posteriors, [[4 / 5, 1 / 5], [1 / 5, 4 / 5]]))
-        else:
-            assert np.all(posteriors == likelihoods)
-    elif prior == [1 / 3, 2 / 3]:
-        if normalize:
-            assert np.all(np.isclose(posteriors, [[2 / 3, 1 / 3], [1 / 9, 8 / 9]]))
-        else:
-            assert np.all(np.isclose(posteriors, [[4 / 3, 2 / 3], [1 / 3, 8 / 3]]))
-    else:
-        raise ValueError("nooo")
+    # Using np.allclose for floating-point precision comparison
+    assert np.allclose(posteriors, expected)
 
 
 def test_posterior_prob_2d_or_mixed_prior():
@@ -147,27 +142,19 @@ def test_posterior_prob_2d_or_mixed_prior():
         _ = classify.posterior_prob(likelihoods, axis=1, normalize=[True, True, False])
 
 
-@pytest.mark.parametrize("likelihoods", (np.array([4, 1]),))
-@pytest.mark.parametrize("prior", (None, [1 / 2, 1 / 2], [1 / 3, 2 / 3]))
-@pytest.mark.parametrize("normalize", (True, False))
-def test_posterior_prob_1d(likelihoods, prior, normalize):
+@pytest.mark.parametrize(
+    "likelihoods, prior, normalize, expected",
+    [
+        (np.array([4, 1]), None, True, [4 / 5, 1 / 5]),
+        (np.array([4, 1]), None, False, [4, 1]),
+        (np.array([4, 1]), [1 / 2, 1 / 2], True, [4 / 5, 1 / 5]),
+        (np.array([4, 1]), [1 / 2, 1 / 2], False, [2, 0.5]),
+        (np.array([4, 1]), [1 / 3, 2 / 3], True, [2 / 3, 1 / 3]),
+        (np.array([4, 1]), [1 / 3, 2 / 3], False, [4 / 3, 2 / 3]),
+    ],
+)
+def test_posterior_prob_1d_(likelihoods, prior, normalize, expected):
     posteriors = classify.posterior_prob(
         likelihoods, axis=0, prior=prior, normalize=normalize
     )
-    if prior == [1 / 2, 1 / 2]:  # hard-coded b/c idk how to engineer tests
-        if normalize:
-            assert np.all(np.isclose(posteriors, [4 / 5, 1 / 5]))
-        else:
-            assert np.all(posteriors == np.array(likelihoods) / 2)
-    elif prior is None:
-        if normalize:
-            assert np.all(np.isclose(posteriors, [4 / 5, 1 / 5]))
-        else:
-            assert np.all(posteriors == likelihoods)
-    elif prior == [1 / 3, 2 / 3]:
-        if normalize:
-            assert np.all(np.isclose(posteriors, [2 / 3, 1 / 3]))
-        else:
-            assert np.all(np.isclose(posteriors, [4 / 3, 2 / 3]))
-    else:
-        raise ValueError("nooo")
+    assert np.allclose(posteriors, expected)