diff --git a/launcher_scripts/nemo_launcher/collections/eval_harness/lm_eval/base.py b/launcher_scripts/nemo_launcher/collections/eval_harness/lm_eval/base.py index f23fa6907..a5f848b68 100755 --- a/launcher_scripts/nemo_launcher/collections/eval_harness/lm_eval/base.py +++ b/launcher_scripts/nemo_launcher/collections/eval_harness/lm_eval/base.py @@ -23,7 +23,16 @@ import numpy as np import spacy from lm_eval.metrics import mean, weighted_mean, weighted_perplexity -from sqlitedict import SqliteDict + +try: + from sqlitedict import SqliteDict + + HAS_SQLITEDICT = True +except ImportError: + print( + "Eval harness with sqlitedict is deprecated. Sqlitedict has known vulnerability GHSA-g4r7-86gm-pgqc" + ) + HAS_SQLITEDICT = False def _SPACY_NLP(*args, **kwargs): @@ -607,7 +616,10 @@ def __init__(self, lm, cache_db): self.cache_db = cache_db if os.path.dirname(cache_db): os.makedirs(os.path.dirname(cache_db), exist_ok=True) - self.dbdict = SqliteDict(cache_db, autocommit=True) + if HAS_SQLITEDICT: + self.dbdict = SqliteDict(cache_db, autocommit=True) + else: + self.dbdict = None # add hook to lm lm.set_cache_hook(self.get_cache_hook()) diff --git a/requirements.txt b/requirements.txt index 84f9b3c7b..2b219f6d0 100755 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,6 @@ pytablewriter==0.58.0 requests==2.26.0 tqdm==4.62.3 zstandard==0.15.2 -sqlitedict hera pydantic kubeflow-training>=1.8