pre-commit

benchmarks/transformers/prob_model_text_classification.py

@@ -36,6 +36,7 @@
  accuracy,
  expected_calibration_error,
 )
+from fortuna.model_editor import ProbitModelEditor
 from fortuna.prob_model import (
  ADVIPosteriorApproximator,
  DeepEnsemblePosteriorApproximator,
@@ -51,7 +52,6 @@
  SNGPPosteriorApproximator,
  SWAGPosteriorApproximator,
 )
-from fortuna.model_editor import ProbitModelEditor
 from fortuna.prob_model.fit_config.hyperparameters import FitHyperparameters
 from fortuna.prob_model.posterior.posterior_approximations import (
  ADVI_NAME,
@@ -400,7 +400,7 @@ def unpack_model_tar(model_ckpt_path: pathlib.Path) -> pathlib.Path:
  if args.enable_probit_model_editor:
  model_editor = ProbitModelEditor(
  freeze_fun=lambda p, v: True if "classifier" in p else False,
- init_log_var=0.
+ init_log_var=0.0,
  )
 
  ### TRAINING
@@ -498,7 +498,9 @@ def unpack_model_tar(model_ckpt_path: pathlib.Path) -> pathlib.Path:
  )
 
  if args.enable_probit_model_editor:
- logger.info(f"Probit log-variance: {prob_model.posterior.state.get().params['model_editor']['params']['log_var']}")
+ logger.info(
+ f"Probit log-variance: {prob_model.posterior.state.get().params['model_editor']['params']['log_var']}"
+ )
 
  ### IN-D PERFORMANCE
  test_inputs_loader = test_data_loader.to_inputs_loader()

fortuna/model_editor/probit.py

@@ -51,6 +51,6 @@ def __call__(
  top_k=self.top_k,
  memory=self.memory,
  n_final_tokens=self.n_final_tokens,
- stop_gradient=self.stop_gradient
+ stop_gradient=self.stop_gradient,
  )
  return outputs

fortuna/utils/probit.py

@@ -69,7 +69,7 @@ def sequential_probit_scaling(
  top_k: Optional[int] = None,
  memory: Optional[int] = None,
  n_final_tokens: Optional[int] = None,
- stop_gradient: bool = False
+ stop_gradient: bool = False,
 ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, Dict]]:
  params = params.unfreeze()
 
@@ -162,7 +162,9 @@ def _compute_cov(_x, idx):
 
  return vmap(J1J2T_op)(jnp.eye(size)).T
 
- return jnp.where(prev_tau != -1, _compute_cov(x, indices), jnp.empty(block_size))
+ return jnp.where(
+ prev_tau != -1, _compute_cov(x, indices), jnp.empty(block_size)
+ )
 
  init_tau = seq_length - n_final_tokens + 1
 
@@ -201,7 +203,10 @@ def fun(carry, tau):
 
  def get_diagCs(_params):
  old_taus = jnp.concatenate(
- (jnp.zeros(memory - 1, dtype="int32") - 1, jnp.array([init_tau], dtype="int32"))
+ (
+ jnp.zeros(memory - 1, dtype="int32") - 1,
+ jnp.array([init_tau], dtype="int32"),
+ )
  )
  C = compute_cov(old_taus[-1], old_taus[-1])