feat: support vllm lora

Author: terryyz

bigcodebench/generate.py

@@ -127,6 +127,7 @@ def run_codegen(
     split: str,
     subset: str,
     root: str = "bcb_results",
+    lora_path: str = None,
     bs: Optional[int] = None,
     n_samples: int = 1,
     temperature: float = 0.0,
@@ -174,6 +175,7 @@ def run_codegen(
         backend=backend,
         subset=subset,
         split=split,
+        lora_path=lora_path,
         temperature=temperature,
         max_new_tokens=max_new_tokens,
         reasoning_effort=reasoning_effort,

bigcodebench/provider/__init__.py

@@ -6,6 +6,7 @@ def make_model(
     backend: str,
     subset: str,
     split: str,
+    lora_path: str = None,
     dataset: str = "bigcodebench",
     temperature: float = 0.0,
     max_new_tokens: int = 1280,
@@ -38,6 +39,7 @@ def make_model(
             name=model,
             subset=subset,
             split=split,
+            lora_path=lora_path,
             temperature=temperature,
             max_new_tokens=max_new_tokens,
             revision=revision,
@@ -58,6 +60,7 @@ def make_model(
             name=model,
             subset=subset,
             split=split,
+            lora_path=lora_path,
             temperature=temperature,
             max_new_tokens=max_new_tokens,
             revision=revision,

bigcodebench/provider/vllm.py

@@ -3,6 +3,8 @@
 
 from transformers import AutoTokenizer
 from vllm import LLM, SamplingParams
+from vllm.lora.request import LoRARequest
+from huggingface_hub import snapshot_download
 
 from bigcodebench.provider.base import DecoderBase
 from bigcodebench.provider.utility import (
@@ -11,7 +13,7 @@
 )
 
 class VllmDecoder(DecoderBase):
-    def __init__(self, name: str, dataset: str, tp: int, **kwargs) -> None:
+    def __init__(self, name: str, lora_path: str, dataset: str, tp: int, **kwargs) -> None:
         super().__init__(name, **kwargs)
 
         kwargs = {
@@ -29,7 +31,17 @@ def __init__(self, name: str, dataset: str, tp: int, **kwargs) -> None:
         else:
             if self.prefill and "```" in self.response_prefix:
                 self.eos += ["\n```\n"]
-        self.llm = LLM(model=name, max_model_len=self.max_new_tokens, **kwargs)
+        
+        self.lora_request = None
+        if lora_path:
+            local_lora_path = snapshot_download(lora_path)
+            self.lora_request = LoRARequest(
+                "lora",
+                1,
+                local_lora_path,
+            )
+        
+        self.llm = LLM(model=name, max_model_len=self.max_new_tokens, enable_lora=True if self.lora_path else False, **kwargs)
         self.llm.set_tokenizer(tokenizer=self.tokenizer)
 
     def is_direct_completion(self) -> bool:
@@ -64,6 +76,7 @@ def codegen(
                 stop=self.eos,
                 skip_special_tokens=self.skip_special_tokens,
             ),
+            lora_request=self.lora_request,
             use_tqdm=True,
         )