vrr/heretic
2
0
Fork 0
mirror of https://github.com/p-e-w/heretic.git synced 2026-07-09 17:28:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Abort trial early if KL divergence is too high

Browse source
This commit is contained in:
Philipp Emanuel Weidmann 2025-09-23 13:20:31 +05:30
parent 9485edc221
commit b6c715ab6f
No known key found for this signature in database
GPG key ID: CE6A7756AD520E0F
2 changed files with 9 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

8
src/heretic/evaluator.py
View file

@ -53,7 +53,13 @@ class Evaluator:
kl_divergence = F.kl_div(
logprobs, self.base_logprobs, reduction="batchmean", log_target=True
).item()
print(f" * KL divergence: [bold]{kl_divergence:.4f}[/]")
print(f" * KL divergence: [bold]{kl_divergence:.4f}[/]", end="")
if kl_divergence > self.settings.max_kl_divergence:
print(" [yellow](constraint violation; aborting trial)[/]")
return -1, kl_divergence, self.base_refusals
else:
print()
print(" * Counting model refusals...")
refusals = self.count_refusals()

2
src/heretic/main.py
View file

@ -198,6 +198,8 @@ def main():
)
print(f" * Score: [bold]{-study.best_value:.4f}[/]")
return
print()
action = questionary.select(
"What do you want to do with the optimized model?",