load tokenizer method (#30)

make all tokenizers to be left padded

Author: mayank31398

inference_server/model_handler/deployment.py

@@ -13,7 +13,7 @@
 from transformers import AutoTokenizer
 
 from ..constants import DS_INFERENCE, DS_ZERO
-from ..models import get_downloaded_model_path, get_model_class
+from ..models import get_downloaded_model_path, get_model_class, load_tokenizer
 from ..utils import (
     GenerateResponse,
     TokenizeRequest,
@@ -32,7 +32,7 @@ def __init__(self, args: argparse.Namespace, use_grpc_server: bool = False, cuda
         self.use_grpc_server = use_grpc_server
 
         if self.use_grpc_server:
-            self.tokenizer = AutoTokenizer.from_pretrained(get_downloaded_model_path(args.model_name))
+            self.tokenizer = load_tokenizer(get_downloaded_model_path(args.model_name))
 
             self.initialize_ports()
 

inference_server/models/__init__.py

@@ -1,5 +1,5 @@
 from ..constants import DS_INFERENCE, DS_ZERO, HF_ACCELERATE
-from .model import Model, get_downloaded_model_path
+from .model import Model, get_downloaded_model_path, load_tokenizer
 
 
 def get_model_class(deployment_framework: str):

inference_server/models/ds_inference.py

@@ -12,7 +12,7 @@
 from transformers import AutoConfig, AutoTokenizer
 
 from ..utils import print_rank_n, run_rank_n
-from .model import Model, get_downloaded_model_path, get_hf_model_class
+from .model import Model, get_downloaded_model_path, get_hf_model_class, load_tokenizer
 
 
 # basic DeepSpeed inference model class for benchmarking
@@ -26,7 +26,7 @@ def __init__(self, args: Namespace) -> None:
 
         downloaded_model_path = get_downloaded_model_path(args.model_name)
 
-        self.tokenizer = AutoTokenizer.from_pretrained(downloaded_model_path)
+        self.tokenizer = load_tokenizer(downloaded_model_path)
         self.pad = self.tokenizer.pad_token_id
 
         # create dummy tensors for allocating space which will be filled with

inference_server/models/ds_zero.py

@@ -9,7 +9,7 @@
 from transformers.deepspeed import HfDeepSpeedConfig
 
 from ..utils import print_rank_n
-from .model import Model, get_downloaded_model_path, get_hf_model_class
+from .model import Model, get_downloaded_model_path, get_hf_model_class, load_tokenizer
 
 
 class DSZeROModel(Model):
@@ -54,7 +54,7 @@ def __init__(self, args: Namespace) -> None:
         # this tells from_pretrained to instantiate directly on gpus
         dschf = HfDeepSpeedConfig(ds_config)
 
-        self.tokenizer = AutoTokenizer.from_pretrained(downloaded_model_path)
+        self.tokenizer = load_tokenizer(downloaded_model_path)
         self.pad = self.tokenizer.pad_token_id
 
         self.model = get_hf_model_class(args.model_class).from_pretrained(

inference_server/models/hf_accelerate.py

@@ -5,7 +5,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from ..utils import print_rank_n
-from .model import Model, get_downloaded_model_path, get_hf_model_class
+from .model import Model, get_downloaded_model_path, get_hf_model_class, load_tokenizer
 
 
 class HFAccelerateModel(Model):
@@ -16,7 +16,7 @@ def __init__(self, args: Namespace) -> None:
 
         downloaded_model_path = get_downloaded_model_path(args.model_name)
 
-        self.tokenizer = AutoTokenizer.from_pretrained(downloaded_model_path)
+        self.tokenizer = load_tokenizer(downloaded_model_path)
         self.pad = self.tokenizer.pad_token_id
 
         kwargs = {"pretrained_model_name_or_path": downloaded_model_path, "device_map": "auto"}

inference_server/models/model.py

@@ -7,7 +7,7 @@
 
 import transformers
 from huggingface_hub import snapshot_download
-from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM
+from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
 from transformers.utils import is_offline_mode
 
 from ..utils import GenerateRequest, GenerateResponse, GenerationMixin, TokenizeRequest, TokenizeResponse, run_rank_n
@@ -115,3 +115,13 @@ def check_batch_size(batch_size: int, max_batch_size: int) -> None:
 # this is a hack for now
 def get_hf_model_class(model_class: str) -> Union[AutoModelForCausalLM, AutoModelForSeq2SeqLM]:
     return getattr(transformers, model_class)
+
+
+def load_tokenizer(model_name: str) -> AutoTokenizer:
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    if tokenizer.pad_token_id is None:
+        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+
+    tokenizer.padding_side = "left"
+    return tokenizer