Skip to content

Commit

Permalink
Make BLEURT lazy (#536)
Browse files Browse the repository at this point in the history
* make bleur lazy

* make tokenizer lazy too
  • Loading branch information
hynky1999 authored Feb 6, 2025
1 parent 1ce7331 commit 15bdbb8
Showing 1 changed file with 15 additions and 3 deletions.
18 changes: 15 additions & 3 deletions src/lighteval/metrics/metrics_sample.py
Original file line number Diff line number Diff line change
Expand Up @@ -708,9 +708,21 @@ def __init__(self):
"""Creates a BLEURT scorer using a light bleurt-tiny-512 model.
For more complex use cases, could also be Elron/bleurt-base-128
"""
self.tokenizer = AutoTokenizer.from_pretrained("Elron/bleurt-tiny-512")
self.model = AutoModelForSequenceClassification.from_pretrained("Elron/bleurt-tiny-512")
self.model.eval()
self._tokenizer = None
self._model = None

@property
def tokenizer(self):
if self._tokenizer is None:
self._tokenizer = AutoTokenizer.from_pretrained("Elron/bleurt-tiny-512")
return self._tokenizer

@property
def model(self):
if self._model is None:
self._model = AutoModelForSequenceClassification.from_pretrained("Elron/bleurt-tiny-512")
self._model.eval()
return self._model

def compute(self, golds: list[str], predictions: list[str], **kwargs) -> float:
"""Uses the stored BLEURT scorer to compute the score on the current sample.
Expand Down

0 comments on commit 15bdbb8

Please sign in to comment.