Optionally return platt scaling classifier.

p-lambda · Nov 8, 2022 · bdd60a7 · bdd60a7
1 parent 48c3dcf
commit bdd60a7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/calibration/utils.py b/calibration/utils.py
@@ -453,7 +453,7 @@ def bootstrap_std(data: List[T], estimator=None, num_samples=100) -> Tuple[float
 
 # Re-Calibration utilities.
 
-def get_platt_scaler(model_probs, labels):
+def get_platt_scaler(model_probs, labels, get_clf=False):
     clf = LogisticRegression(C=1e10, solver='lbfgs')
     eps = 1e-12
     model_probs = model_probs.astype(dtype=np.float64)
@@ -468,6 +468,8 @@ def calibrator(probs):
         x = x * clf.coef_[0] + clf.intercept_
         output = 1 / (1 + np.exp(-x))
         return output
+    if get_clf:
+        return calibrator, clf
     return calibrator
 
 

diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="uncertainty-calibration",
-    version="0.1.2",
+    version="0.1.3",
     author="Ananya Kumar",
     author_email="[email protected]",
     description="Utilities to calibrate model uncertainties and measure calibration.",