sentence_piece.py

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-4.0

"""
Script for training, encoding, and decoding sub-word segmentation model
using SentencePiece.
"""

import sentencepiece as sp
import argparse
from contextlib import ExitStack

def train(sp_model, inputs, vocab_size, model_type):
  sp.SentencePieceTrainer.Train('--input={} --model_prefix={} --vocab_size={} '
          '--model_type={}'.format(inputs, sp_model, vocab_size, model_type))

def encode_decode(sp_model, inputs, outputs, mode=None):
  sp_spp = sp.SentencePieceProcessor()
  sp_spp.Load(sp_model)
  
  if mode == "encode":
    def preprocess(sample):
        return " ".join(sp_spp.EncodeAsPieces(sample.strip()))
  if mode == "decode":
    def preprocess(sample):
        return "".join(sp_spp.DecodePieces(sample.strip().split()))

  with ExitStack() as stack:
      infiles = [stack.enter_context(open(infile, "r", encoding="utf-8")) for infile in inputs]
      outfiles = [stack.enter_context(open(outfile, "w", encoding="utf-8")) for outfile in outputs]

      for s, samples in enumerate(zip(*infiles)):
        processed_samples = list(map(preprocess, samples))
        for sample, outfile in zip(processed_samples, outfiles):
          print(sample, file=outfile)

def main():
  parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument("--run", choices=['train', 'encode', 'decode'], required=True,
                      help="SP run type.")
  parser.add_argument("--sp_model", required=True, 
                      help="SP model." )
  parser.add_argument("--inputs", nargs="+", type=str, default="",
                      help="input files for SP run type.")
  parser.add_argument("--vocab_size", required=False, default=3200, 
                      help="number of SP segmentation rules.")
  parser.add_argument("--model_type", required=False, default="bpe",
                      help="SP model type (bpe, unigram)." )
  parser.add_argument("--outputs", nargs="+", required=False, default="", 
                      help="files for SP run outputs.")
  args = parser.parse_args()

  if args.run == "train":
    inputs = ",".join(args.inputs)
    print(f"SentencePiece: {args.run} ...")
    train(args.sp_model, inputs, args.vocab_size, args.model_type)  
  elif args.run == "encode" or args.run == "decode":
    assert len(args.inputs) == len(args.outputs), "different input and output file size."
    print(f"SentencePiece: {args.run} ...")
    encode_decode(args.sp_model, args.inputs,args. outputs, args.run)
  else:
    raise ValueError(f"SP run mode {args.run} unknown.")

if __name__ == "__main__":
  main()