feat: add VLLM wrapper

config/llm_template.json

@@ -7,6 +7,10 @@
         "template": "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ '<s>' + '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' '  + content.strip() + ' ' + '</s>' }}{% endif %}{% endfor %}",
         "system_prompt": true
     },
+    "llama-3": {
+        "template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = '<|begin_of_text|>' + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
+        "system_prompt": false
+    },
     "mistral": {
         "template": "'<s>'+{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content']}}+'</s>'+{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}",
         "system_prompt": true

env.template

@@ -1,7 +1,11 @@
 #GPT
-AZURE_KEY=xxxx123
-AZURE_VERSION=2023-07-01-preview
-AZURE_ENDPOINT=http://gpt-4.com
+OPENAI_API_TYPE="azure"
+OPENAI_API_BASE="https://<your-endpoint.openai.azure.com/>"
+OPENAI_API_KEY="your AzureOpenAI key"
+OPENAI_API_VERSION="2023-05-15"
+
+#TGI
+TGI_ENDPOINT=""
 
 #GEMINI
-GEMINI_KEY=abcdefghtyuidkg
+GEMINI_KEY="abcdefghtyuidkg"

pyproject.toml

@@ -29,6 +29,7 @@ classifiers = [
 license = { text = "Apache 2.0 License" }
 
 dependencies = [
+  "vllm>=0.5.2",
   "accelerate>=0.30.1",
   "peft>=0.11.1",
   "bitsandbytes>=0.40.2",

src/vieval/script_arguments.py

@@ -90,7 +90,9 @@ class ScriptArguments:
         default=True,
         metadata={"help": "Enable the TF32 mode (available in Ampere and newer GPUs)"},
     )
-
+    cpu_offload_gb: Optional[bool] = field(
+        default=0, metadata={"help": "Offload computation to CPU"}
+    )
     auto_find_batch_size: Optional[bool] = field(
         default=True, metadata={"help": "Enable auto batch size"}
     )
@@ -167,6 +169,9 @@ class ScriptArguments:
     )
 
     # Inference parameters
+    dtype: Optional[str] = field(
+        default="float16", metadata={"help": "Data type for model weights"}  # float16
+    )
     smoke_test: Optional[bool] = field(
         default=False, metadata={"help": "Run a smoke test on a small dataset"}
     )
@@ -182,7 +187,6 @@ class ScriptArguments:
     cot: Optional[bool] = field(
         default=False, metadata={"help": "Enable chain of thought when prompting MATH"}
     )
-    tgi: Optional[str] = field(default="", metadata={"help": "Embed TGI endpoint"})
     seed: Optional[int] = field(default=42, metadata={"help": "Random seed"})
     continue_infer: Optional[bool] = field(
         default=False,

src/vieval/tools/pipelines/pipelines.py

@@ -4,7 +4,7 @@
 import json
 from tqdm import tqdm
 from ..utils.model import get_model
-from ..wrapper import AzureGPTWrapper, TGIWrapper, GeminiWrapper, HFWrapper
+from ..wrapper import OpenAIWrapper, TGIWrapper, GeminiWrapper, HFWrapper
 from ..utils.utils import *
 from ..utils.metric_utils import info_from_filename
 from .metric_pipelines import MetricPipeline
@@ -28,24 +28,23 @@ def __init__(self, task, config):
         # print(config.tgi)
         if config.wtype == "tgi":
             self.infer_pipeline = TGIWrapper(
-                api_endpoint=config.tgi,
                 generation_config=GenerationConfig[extract_task],
                 template=LLM_TEMPLATE[config.ptemplate],
             )
         elif config.wtype == "hf":
-            # Load model
-            self.model, self.tokenizer = get_model(config=config)
-            self.model.eval()
-
             self.infer_pipeline = HFWrapper(
-                model=self.model,
-                tokenizer=self.tokenizer,
+                config=config,
                 generation_config=GenerationConfig[extract_task],
                 template=LLM_TEMPLATE[config.ptemplate],
             )
-
-        elif config.wtype == "azuregpt":
-            self.infer_pipeline = AzureGPTWrapper(
+        elif config.wtype == "vllm":
+            self.infer_pipeline = VLLMWrapper(
+                config=config,
+                generation_config=GenerationConfig[extract_task],
+                template=LLM_TEMPLATE[config.ptemplate],
+            )
+        elif config.wtype == "openai":
+            self.infer_pipeline = OpenAIWrapper(
                 engine=config.model_name,
                 generation_config=GenerationConfig[extract_task],
             )
@@ -146,7 +145,7 @@ def __question_answering(self, ds_wrapper, ds_loader, saving_fn, start_idx=0):
             results, logprobs, _ = self.infer_pipeline(prompts, return_probs=True)
             predictions.extend(results)
             references.extend([x[0] for x in batch[ds_wrapper.answer]["text"]])
-            generation_probs.extend([x.tolist() for x in logprobs])
+            generation_probs.extend(logprobs)
 
             idx += 1
             if idx % 100 == 0:
@@ -185,7 +184,7 @@ def __question_answering_without_context(
         calib_probs = []
         idx = 0
         original_few_shot = []
-        calibration_few_shot = []
+        calib_few_shot = []
         selected_sample = []
         if self.continue_infer_data is not None:
             predictions.extend(self.continue_infer_data["predictions"])
@@ -258,8 +257,8 @@ def preprocessing_a_record(rec):
             )
             predictions.extend(results)
             references.extend([x for x in batch[ds_wrapper.answer]])
-            generation_probs.extend([x.tolist() for x in logprobs])
-            calib_probs.extend([x.tolist() for x in calibprob_batch])
+            generation_probs.extend(logprobs)
+            calib_probs.extend(calibprob_batch)
             idx += 1
             if idx % 100 == 0:
                 print(f"Saving results of {idx} batches")
@@ -327,7 +326,7 @@ def __summarization(self, ds_wrapper, ds_loader, saving_fn, start_idx=0):
             results, logprobs, _ = self.infer_pipeline(prompts, return_probs=True)
             predictions.extend(results)
             references.extend([x for x in batch[ds_wrapper.summarized_text]])
-            generation_probs.extend([x.tolist() for x in logprobs])
+            generation_probs.extend(logprobs)
 
             idx += 1
             if idx % 100 == 0:
@@ -451,7 +450,7 @@ def preprocessing_a_record(rec):
             )
             predictions.extend(results)
             references.extend([x.item() for x in batch[ds_wrapper.label]])
-            generation_probs.extend([x.tolist() for x in logprobs])
+            generation_probs.extend(logprobs)
             option_probs.extend(
                 [
                     [
@@ -623,7 +622,7 @@ def format_calib_fewshot(rec):
                     for x in batch[ds_wrapper.label]
                 ]
             )
-            generation_probs.extend([x.tolist() for x in logprobs])
+            generation_probs.extend(logprobs)
             option_probs.extend(
                 [
                     [
@@ -756,7 +755,7 @@ def preprocessing_a_record(rec):
             )
             predictions.extend(results)
             references.extend([x.item() for x in batch[ds_wrapper.label]])
-            generation_probs.extend([x.tolist() for x in logprobs])
+            generation_probs.extend(logprobs)
             option_probs.extend(
                 [
                     [
@@ -932,7 +931,7 @@ def preprocessing_a_record(rec):
                 ]
             )
 
-            generation_probs.extend([x.tolist() for x in logprobs])
+            generation_probs.extend(logprobs)
             option_probs.extend(opt_calib_out)
             idx += 1
             if idx % 100 == 0:
@@ -1017,7 +1016,7 @@ def preprocessing_a_record(rec):
             results, logprobs, _ = self.infer_pipeline(prompts, return_probs=True)
             predictions.extend(results)
             references.extend([x for x in batch[ds_wrapper.target]])
-            generation_probs.extend([x.tolist() for x in logprobs])
+            generation_probs.extend(logprobs)
 
             idx += 1
             if idx % 100 == 0:
@@ -1322,8 +1321,8 @@ def preprocessing_a_record(rec):
             )
             predictions.extend(results)
             references.extend([x for x in batch[target]])
-            generation_probs.extend([x.tolist() for x in logprobs])
-            calib_probs.extend([x.tolist() for x in calibprob_batch])
+            generation_probs.extend(logprobs)
+            calib_probs.extend(calibprob_batch)
             if sub_task == "math":
                 math_problem_type.extend([x for x in batch[ds_wrapper.type]])
             idx += 1
@@ -1417,7 +1416,7 @@ def preprocessing_a_record(rec):
             results, logprobs, _ = self.infer_pipeline(prompts, return_probs=True)
             predictions.extend(results)
             references.extend([x for x in batch[ds_wrapper.target_language]])
-            generation_probs.extend([x.tolist() for x in logprobs])
+            generation_probs.extend(logprobs)
 
             idx += 1
             if idx % 100 == 0:

src/vieval/tools/wrapper/HFWrapper.py

@@ -1,12 +1,14 @@
 import torch
 from .BaseWrapper import BaseWrapper
 from ..utils.chat_template import apply_chat_template
+from ..utils.model import get_model
 
 
 class HFWrapper(BaseWrapper):
-    def __init__(self, model, tokenizer, generation_config, template=""):
-        self.model = model
-        self.tokenizer = tokenizer
+    def __init__(self, config, generation_config, template=None):
+        self.model, self.tokenizer = get_model(config=config)
+        self.model.eval()
+
         self.generation_config = generation_config
         self.model_template = template
 
@@ -51,7 +53,7 @@ def __call__(self, prompts, return_probs=False):
                     scores=generate_dict.scores,
                     normalize_logits=True,
                 )
-                generations_probs.extend(generation_probs.cpu().numpy())
+                generations_probs.extend(generation_probs.cpu().numpy().tolist())
 
         return generations, generations_probs, num_generated_tokens
 
@@ -67,7 +69,7 @@ def compute_logprob_and_length(self, prompts, completions):
             prompt_num_tokens = prompt_tokens.input_ids.shape[1] - 1
 
             completion_tokens = self.tokenizer(
-                f"{completion} {self.tokenizer.eos_token}", return_tensors="pt"
+                f"{completion}{self.tokenizer.eos_token}", return_tensors="pt"
             ).to(
                 self.model.device
             )  # <s> SPIECE_UNDERLINE [tokens] SPIECE_UNDERLINE </s>
@@ -101,5 +103,5 @@ def compute_logprob_and_length(self, prompts, completions):
                 ),
             ).squeeze(-1)
             # >>> batch_size, sequence_length
-            completions_logprobs.append(logprobs.cpu().numpy())
+            completions_logprobs.append(logprobs.cpu().numpy().tolist())
         return completions_logprobs, completions_num_tokens

src/vieval/tools/wrapper/AzureGPTWrapper.py → src/vieval/tools/wrapper/OpenAIWrapper.py

@@ -5,14 +5,14 @@
 from .BaseWrapper import BaseWrapper
 
 
-class AzureGPTWrapper(BaseWrapper):
+class OpenAIWrapper(BaseWrapper):
     def __init__(self, engine=None, generation_config=None):
-        self.generation_config = generation_config
-        self.model = openai.AzureOpenAI(
-            azure_endpoint=os.getenv("AZURE_ENDPOINT"),
-            api_key=os.getenv("AZURE_KEY"),
-            api_version=os.getenv("AZURE_VERSION"),
+        generation_config["max_tokens"] = generation_config.pop("max_new_tokens")
+        generation_config["frequency_penalty"] = generation_config.pop(
+            "repetition_penalty"
         )
+        self.generation_config = generation_config
+        self.model = openai.OpenAI()
         self.engine = engine
 
     def __call__(self, prompts, return_probs=False):
@@ -24,10 +24,7 @@ def __call__(self, prompts, return_probs=False):
             response = self.chat_completions_with_backoff(
                 model=self.engine,
                 messages=prompt,
-                temperature=self.generation_config["temperature"],
-                max_tokens=self.generation_config["max_new_tokens"],
-                top_p=0.95,
-                frequency_penalty=self.generation_config["repetition_penalty"],
+                **self.generation_config,
             )
 
             generations.append(response.choices[0].message.content)