Fix Transformers bug in NeMo Curator (#591)

Signed-off-by: Vibhu Jawa <vjawa@nvidia.com>

Author: VibhuJawa

nemo_curator/classifiers/aegis.py

@@ -24,7 +24,6 @@
 from crossfit import op
 from crossfit.backend.torch.hf.model import HFModel
 from huggingface_hub import PyTorchModelHubMixin
-from peft import PeftModel
 from torch.nn import Dropout, Linear
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
@@ -116,6 +115,11 @@ def __init__(
         base_model = AutoModelForCausalLM.from_pretrained(
             pretrained_model_name_or_path, torch_dtype=dtype, token=token
         )
+        # Importing PeftModel here to prevent cuda context issues
+        # that seem to happen on Transformers 4.48.3
+        # See related: https://github.com/rapidsai/crossfit/pull/113
+        from peft import PeftModel
+
         self.model = PeftModel.from_pretrained(base_model, peft_model_name_or_path)
         self.autocast = autocast
         self.add_instruction_data_guard = add_instruction_data_guard

pyproject.toml

@@ -43,7 +43,7 @@ dependencies = [
     "beautifulsoup4",
     "charset_normalizer>=3.1.0",
     "comment_parser",
-    "crossfit>=0.0.8",
+    "crossfit>=0.0.8.post1",
     "dask-mpi>=2021.11.0",
     "dask[complete]>=2021.7.1",
     "datasets",