Skip to content

Commit

Permalink
Add tokenizer backend (EleutherAI#1186)
Browse files Browse the repository at this point in the history
* separate local flag

* tokenizer_backend

* import order
  • Loading branch information
anjor authored Dec 21, 2023
1 parent 2b0b6fd commit a0cfe3f
Showing 1 changed file with 9 additions and 3 deletions.
12 changes: 9 additions & 3 deletions lm_eval/models/openai_completions.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import time
from collections import defaultdict
from importlib.util import find_spec
from typing import List, Optional, Tuple
from typing import List, Literal, Optional, Tuple

import transformers
from tqdm import tqdm
Expand Down Expand Up @@ -360,6 +360,7 @@ def __init__(
self,
model: str = "gpt-3.5-turbo", # GPT model or Local model using HuggingFace model paths
base_url: str = None,
tokenizer_backend: Literal["tiktoken", "huggingface"] = "tiktoken",
truncate: bool = False,
revision: Optional[str] = "main",
trust_remote_code: Optional[bool] = False,
Expand Down Expand Up @@ -388,10 +389,11 @@ def __init__(
)
self.model = model
self.base_url = base_url
self.tokenizer_backend = tokenizer_backend
self.truncate = truncate

# if we have a local model, use HF tokenizer over tiktoken
if self.base_url:
if self.tokenizer_backend == "huggingface":
self.revision = revision
self.trust_remote_code = trust_remote_code
self.use_fast_tokenizer = use_fast_tokenizer
Expand All @@ -404,10 +406,14 @@ def __init__(
)
self.vocab_size = self.tokenizer.vocab
self.end_of_text_token_id = self.tokenizer.eos_token
else:
elif self.tokenizer_backend == "tiktoken":
self.tokenizer = tiktoken.encoding_for_model(self.model)
self.vocab_size = self.tokenizer.n_vocab
self.end_of_text_token_id = self.tokenizer.eot_token
else:
raise ValueError(
f"Expected tokenizer_backend to be one of ['tiktoken', 'huggingface'] but got {self.tokenizer_backend}"
)

# Read from environment variable OPENAI_API_KEY
# Set to EMPTY for local
Expand Down

0 comments on commit a0cfe3f

Please sign in to comment.