Avoid peft import in global scope, loads bitsandbytes and that loads …

…cuda and corrupts parent
h2oai · Jun 9, 2023 · 762cdcd · 762cdcd
1 parent 897a0a8
commit 762cdcd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/generate.py b/generate.py
@@ -33,7 +33,6 @@
 
 import fire
 import torch
-from peft import PeftModel
 from transformers import GenerationConfig, AutoModel, TextIteratorStreamer
 from accelerate import init_empty_weights, infer_auto_device_map
 
@@ -710,6 +709,7 @@ def get_model(
                 base_model,
                 **model_kwargs
             )
+            from peft import PeftModel  # loads cuda, so avoid in global scope
             model = PeftModel.from_pretrained(
                 model,
                 lora_weights,
@@ -727,6 +727,7 @@ def get_model(
                     base_model,
                     **model_kwargs
                 )
+                from peft import PeftModel  # loads cuda, so avoid in global scope
                 model = PeftModel.from_pretrained(
                     model,
                     lora_weights,

diff --git a/utils.py b/utils.py
@@ -14,7 +14,6 @@
 import traceback
 import zipfile
 from datetime import datetime
-from enum import Enum
 
 import filelock
 import requests, uuid