Skip to content

ci: add guards for lm eval imports to make importing the metric okay #610

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gsprochette merged 2 commits into from
Apr 3, 2026
+12 −6
Merged

ci: add guards for lm eval imports to make importing the metric okay #610

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b28ef12
ci: add guards for lm eval imports to make importing the metric okay
begumcig Apr 2, 2026
c3a01a3
style: add back comment explaining lm_eval.api.metrics import and sli…
gsprochette Apr 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 12 additions & 6 deletions src/pruna/evaluation/metrics/metric_evalharness.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,15 @@
from typing import Any, List, Tuple

import torch
from lm_eval.api import metrics # noqa: F401 # needed to register lm-eval metrics
from lm_eval.api import registry as lm_registry

try:
# lm_eval.api.metrics import is needed to register lm-eval metrics
from lm_eval.api import metrics as _lm_metrics # noqa: F401
from lm_eval.api import registry as lm_registry

_LM_EVAL_AVAILABLE = True
except ImportError:
_LM_EVAL_AVAILABLE = False

from pruna.evaluation.metrics.metric_stateful import StatefulMetric
from pruna.evaluation.metrics.registry import MetricRegistry
Expand All @@ -29,10 +36,7 @@
METRIC_EVALHARNESS = "lm_eval_metric"
# Perplexity from lm_eval requires the call type to be "y"
# So requires us to handle turning the vector logits into a single log prob.
LM_EVAL_METRICS = [
m for m in lm_registry.METRIC_REGISTRY
if m != "perplexity" # We use the torchmetrics implementation for perplexity
]
LM_EVAL_METRICS = [m for m in lm_registry.METRIC_REGISTRY if m != "perplexity"] if _LM_EVAL_AVAILABLE else []


@MetricRegistry.register_wrapper(available_metrics=LM_EVAL_METRICS)
Expand All @@ -54,6 +58,8 @@ class LMEvalMetric(StatefulMetric):
pairs: List[Tuple[Any, Any]] # dynamically added by add_state()

def __init__(self, metric_name: str, call_type: str = "y_gt") -> None:
if not _LM_EVAL_AVAILABLE:
raise ImportError("lm-eval is required for LMEvalMetric. Install it with: pip install 'pruna[lmharness]'")
super().__init__()
self.metric_name = metric_name
self.call_type = call_type
Expand Down
Loading