From 0c43782c6c31a2339e56b1ee6a14bf27816a39f4 Mon Sep 17 00:00:00 2001 From: Allen Goodman Date: Wed, 22 May 2024 10:24:52 -0400 Subject: [PATCH] cleanup --- .../_protein_mlm_tokenizer_transform.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/beignet/transforms/_protein_mlm_tokenizer_transform.py b/src/beignet/transforms/_protein_mlm_tokenizer_transform.py index ebf563d5e8..00a24d5b28 100644 --- a/src/beignet/transforms/_protein_mlm_tokenizer_transform.py +++ b/src/beignet/transforms/_protein_mlm_tokenizer_transform.py @@ -1,6 +1,6 @@ import importlib.resources from os import PathLike -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Union from transformers.tokenization_utils_base import ( BatchEncoding, @@ -16,24 +16,24 @@ class ProteinMLMTokenizerTransform(Transform): def __init__( self, - pretrained_model_name_or_path: Union[str, PathLike] = None, - padding: Union[bool, str, PaddingStrategy] = False, - truncation: Union[bool, str, TruncationStrategy] = False, - max_length: Optional[int] = None, - return_token_type_ids: Optional[bool] = None, - return_attention_mask: Optional[bool] = None, + path: str | PathLike = None, + padding: bool | str | PaddingStrategy = False, + truncation: bool | str | TruncationStrategy = False, + max_length: int | None = None, + return_token_type_ids: bool | None = None, + return_attention_mask: bool | None = None, return_overflowing_tokens: bool = False, return_special_tokens_mask: bool = False, return_offsets_mapping: bool = False, return_length: bool = False, verbose: bool = True, - tokenizer_dir: Optional[str] = "pmlm_tokenizer", + tokenizer_dir: str | None = "pmlm_tokenizer", mlm: bool = True, reversal_augmentation: bool = False, ): super().__init__() - self._pretrained_model_name_or_path = pretrained_model_name_or_path + self._pretrained_model_name_or_path = path self._padding = padding self._truncation = truncation self._max_length = max_length