Train SentencePiece tokenizer from subset of tf dataset. Args: dataset: tf.dataset vocab_size: int: size of vocab tokens to train. maxchars: int: number of characters to use for sentencepiece training. model_path: str: path of model file to save vocab model to. model_type: str:
(
dataset: tf.data.Dataset,
*,
vocab_size: int,
maxchars: int = int(1e7),
model_path: str,
model_type: str = "unigram",
character_coverage: float = 1.0,
data_keys=("text",),
)
| 64 | |
| 65 | |
| 66 | def _train_sentencepiece( |
| 67 | dataset: tf.data.Dataset, |
| 68 | *, |
| 69 | vocab_size: int, |
| 70 | maxchars: int = int(1e7), |
| 71 | model_path: str, |
| 72 | model_type: str = "unigram", |
| 73 | character_coverage: float = 1.0, |
| 74 | data_keys=("text",), |
| 75 | ): |
| 76 | """Train SentencePiece tokenizer from subset of tf dataset. |
| 77 | Args: |
| 78 | dataset: tf.dataset |
| 79 | vocab_size: int: size of vocab tokens to train. |
| 80 | maxchars: int: number of characters to use for sentencepiece training. |
| 81 | model_path: str: path of model file to save vocab model to. |
| 82 | model_type: str: type of sentencepiece vocab to train. |
| 83 | character_coverage: amount of characters covered by the model, good defaults |
| 84 | are 0.9995 for languages with rich character set like Japanese or Chinese |
| 85 | and 1.0 for other languages with small character set. |
| 86 | data_keys: tuple[str]: keys of dataset to use for training. |
| 87 | Returns: |
| 88 | path to the trained sentencepiece vocabulary model. |
| 89 | """ |
| 90 | if model_path.startswith("gs://"): |
| 91 | abs_model_path = model_path |
| 92 | else: |
| 93 | abs_model_path = os.path.abspath(os.path.expanduser(model_path)) |
| 94 | fname, _ = _dump_chars_to_textfile(dataset, maxchars=maxchars, data_keys=data_keys) |
| 95 | temp_dir = tempfile.gettempdir() |
| 96 | with tempfile.NamedTemporaryFile(delete=False, prefix=os.path.join(temp_dir, "sp_tmp")) as model_fp: |
| 97 | pass # we just want a prefix'd tmp-filename |
| 98 | argstr = " ".join( |
| 99 | [ |
| 100 | f"--input={fname}", |
| 101 | f"--vocab_size={vocab_size}", |
| 102 | f"--character_coverage={character_coverage}", |
| 103 | f"--model_prefix={model_fp.name}", |
| 104 | f"--model_type={model_type}", |
| 105 | ] |
| 106 | ) |
| 107 | SentencePieceTrainer.Train(argstr) |
| 108 | if jax.process_index() == 0: |
| 109 | # Use an intermediate filename that is renamed to the target name to address |
| 110 | # create and fill delays. |
| 111 | copy_rename_path = abs_model_path + ".rntmp" |
| 112 | tf.io.gfile.makedirs(os.path.dirname(abs_model_path)) |
| 113 | tf.io.gfile.copy(model_fp.name + ".model", copy_rename_path, overwrite=True) |
| 114 | tf.io.gfile.rename(copy_rename_path, abs_model_path, overwrite=True) |
| 115 | logging.info("copied %s to %s", model_fp.name + ".model", abs_model_path) |
| 116 | else: |
| 117 | while not tf.io.gfile.exists(abs_model_path): |
| 118 | time.sleep(1) |
| 119 | time.sleep(1) |
| 120 | return abs_model_path |
| 121 | |
| 122 | |
| 123 | def train_tokenizer( |
no test coverage detected