[Python] BLEURT evaluation metric consumed too much RAM

The BLEURT codes almost used up all the 24GB RAM of NVIDIA GeForce RTX 4090 to evaluate just 1 set of sentences.

ref = 'reference sentence here'
hypo = 'hypothesis sentence here'

scores = {}
import evaluate
metric_bleurt = evaluate.load('bleurt')
bleurt = metric_bleurt.compute(predictions=[hypo], references=[ref])
scores["bleurt"] = round(bleurt["scores"][0], 3)
print(scores)


Total memory spent is 22290MB. How is that possible? How do I reduce memory usage when evaluating BLEURT? Thanks.


UPDATE With the evaluate library, the default checkpoint loaded is BLEURT-BASE-128. To set a smaller checkpoint, I use metric_bleurt = evaluate.load('bleurt', config_name="bleurt-tiny-128") (which is undocumented). Still, 22GB RAM is occupied when using BLEURT Tiny checkpoint.

Continue reading...