Skip to content

Commit

Permalink
Merge pull request caikit#309 from gkumbhat/add_model_load_logging
Browse files Browse the repository at this point in the history
🔊 Add model and tokenizer time logging
  • Loading branch information
gkumbhat authored Jan 27, 2024
2 parents 6bc17e2 + f255c4c commit c20c444
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 15 deletions.
1 change: 1 addition & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ repos:
hooks:
- id: black
exclude: imports
additional_dependencies: ["platformdirs"]
- repo: https://github.com/PyCQA/isort
rev: 5.11.5
hooks:
Expand Down
32 changes: 17 additions & 15 deletions caikit_nlp/resources/pretrained_model/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -201,25 +201,27 @@ def bootstrap(
else "right"
)

# Load the tokenizer and set up the pad token if needed
tokenizer = AutoTokenizer.from_pretrained(
tokenizer_name,
local_files_only=not get_config().allow_downloads,
padding_side=padding_side,
# We can't disable use_fast otherwise unit test fails
# use_fast=False,
)
with alog.ContextTimer(log.info, "Tokenizer loaded in "):
# Load the tokenizer and set up the pad token if needed
tokenizer = AutoTokenizer.from_pretrained(
tokenizer_name,
local_files_only=not get_config().allow_downloads,
padding_side=padding_side,
# We can't disable use_fast otherwise unit test fails
# use_fast=False,
)

if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id

# Load the model
model = cls.MODEL_TYPE.from_pretrained(
model_name,
local_files_only=not get_config().allow_downloads,
torch_dtype=torch_dtype,
**kwargs,
)
with alog.ContextTimer(log.info, f"Model {model_name} loaded in "):
# Load the model
model = cls.MODEL_TYPE.from_pretrained(
model_name,
local_files_only=not get_config().allow_downloads,
torch_dtype=torch_dtype,
**kwargs,
)
log.debug4("Model Details: %s", model)

# Create the class instance
Expand Down

0 comments on commit c20c444

Please sign in to comment.