Support loading PEFT (LoRA) models

Author: idoru

basaran/__init__.py

@@ -18,6 +18,7 @@ def is_true(value):
 PORT = int(os.getenv("PORT", "80"))
 
 # Model-related arguments:
+MODEL_PEFT = is_true(os.getenv("MODEL_PEFT", ""))
 MODEL_REVISION = os.getenv("MODEL_REVISION", "")
 MODEL_CACHE_DIR = os.getenv("MODEL_CACHE_DIR", "models")
 MODEL_LOAD_IN_8BIT = is_true(os.getenv("MODEL_LOAD_IN_8BIT", ""))

basaran/__main__.py

@@ -23,6 +23,7 @@
 from . import MODEL_LOAD_IN_4BIT
 from . import MODEL_4BIT_QUANT_TYPE
 from . import MODEL_4BIT_DOUBLE_QUANT
+from . import MODEL_PEFT
 from . import MODEL_LOCAL_FILES_ONLY
 from . import MODEL_TRUST_REMOTE_CODE
 from . import MODEL_HALF_PRECISION
@@ -44,6 +45,7 @@
     name_or_path=MODEL,
     revision=MODEL_REVISION,
     cache_dir=MODEL_CACHE_DIR,
+    is_peft=MODEL_PEFT,
     load_in_8bit=MODEL_LOAD_IN_8BIT,
     load_in_4bit=MODEL_LOAD_IN_4BIT,
     quant_type=MODEL_4BIT_QUANT_TYPE,

basaran/model.py

@@ -14,6 +14,10 @@
     TopPLogitsWarper,
     BitsAndBytesConfig
 )
+from peft import (
+    PeftConfig,
+    PeftModel
+)
 
 from .choice import map_choice
 from .tokenizer import StreamTokenizer
@@ -302,6 +306,7 @@ def load_model(
     name_or_path,
     revision=None,
     cache_dir=None,
+    is_peft=False,
     load_in_8bit=False,
     load_in_4bit=False,
     quant_type="fp4",
@@ -346,6 +351,10 @@ def load_model(
         if half_precision or load_in_8bit or load_in_4bit:
             kwargs["torch_dtype"] = torch.float16
 
+    if is_peft:
+        peft_config = PeftConfig.from_pretrained(name_or_path)
+        name_or_path = peft_config.base_model_name_or_path
+
     # Support both decoder-only and encoder-decoder models.
     try:
         model = AutoModelForCausalLM.from_pretrained(name_or_path, **kwargs)

requirements.txt

@@ -12,3 +12,4 @@ safetensors~=0.3.1
 torch>=1.12.1
 transformers[sentencepiece]~=4.29.2
 waitress~=2.1.2
+peft~=0.3.0