examine + new speak

train/tokenise.py

@@ -0,0 +1,35 @@
+import argparse, os, sys
+
+
+def encode(filepath: str, blocksize: int, ouputdir: str, verbose: bool = False) -> int:
+
+	from aitextgen.TokenDataset import TokenDataset
+	from aitextgen.tokenizers import train_tokenizer
+
+	fn = ouputdir + os.path.basename(filepath)
+
+	#NOTE: vocab_size is fixed since this is not yet in train_tokenizer
+	#see https://github.com/minimaxir/aitextgen/blob/master/aitextgen/tokenizers.py
+
+	train_tokenizer(filepath, prefix=fn)
+
+	tok_fn = fn + ".tokenizer.json"
+	fn_dest = fn + "_bs=" + str(blocksize) + ".tar.gz"
+
+	data = TokenDataset(file_path=filepath, tokenizer_file=tok_fn, block_size=blocksize, line_by_line=True)
+	data.save(cache_destination=fn_dest)
+
+	return 0	
+
+def main() -> int:
+	p = argparse.ArgumentParser()
+	p.add_argument("text", type=str, help="text file path to be tokenised and encoded")
+	p.add_argument("-b", "--blocksize", type=int, choices=[32, 64, 128, 256, 1024], default=64, help="block size, default=64 (corresponds to GPT-2 'max_lenght' config)")
+	p.add_argument("--ouputdir", type=str, default="data/tokens/")
+	p.add_argument("-v", "--verbose", action="store_true")
+	args = p.parse_args()
+
+	return encode(args.text, args.blocksize, args.ouputdir, args.verbose)
+
+if __name__ == '__main__':
+    sys.exit(main())

train/train.py

@@ -0,0 +1,57 @@
+import argparse, os, sys
+from aitextgen.TokenDataset import TokenDataset
+from aitextgen.utils import GPT2ConfigCPU
+from aitextgen.utils import build_gpt2_config
+from aitextgen import aitextgen
+
+
+# https://github.com/minimaxir/aitextgen/blob/master/aitextgen/utils.py
+# https://github.com/huggingface/transformers/blob/master/src/transformers/models/gpt2/configuration_gpt2.py
+
+def run_cpu(te: str, tok: str, dat: str, blocksize: int, num_steps: int = 10000) -> int:
+
+	config = GPT2ConfigCPU()
+
+	ai = aitextgen(tokenizer_file=tok, config=config)
+	data = TokenDataset(dat, tokenizer_file=tok, block_size=blocksize, from_cache=True)
+
+	ai.train(data, output_dir=te, batch_size=16, num_steps=num_steps, generate_every=1000, save_every=1000, num_workers=4)
+
+	return 0
+
+def run_gpu(te: str, tok: str, dat: str, blocksize: int, num_steps: int = 10000) -> int:
+
+	#NOTE: vocab_size is fixed since this is not yet in train_tokenizer
+
+	config = build_gpt2_config(vocab_size=1000, max_lenght=blocksize)
+
+	ai = aitextgen(tokenizer_file=tok, config=config)
+	data = TokenDataset(dat, tokenizer_file=tok, block_size=blocksize, from_cache=True)
+
+	ai.train(data, output_dir=te, batch_size=16, num_steps=num_steps, generate_every=1000, save_every=1000, num_workers=4, to_gpu=True)
+
+	return 0
+
+
+def main() -> int:
+	p = argparse.ArgumentParser()
+	p.add_argument("text", type=str, help="text to create model from")
+	p.add_argument("-b", "--blocksize", type=int, choices=[32, 64, 128, 256, 1024], default=64, help="block size, default=64 (corresponds to GPT-2 'max_lenght' config)")
+	p.add_argument("-s", "--numsteps", type=int, default=10000)
+	p.add_argument("--tokensdir", type=str, default="data/tokens/")
+	p.add_argument("--ouputdir", type=str, default="data/models/")
+	p.add_argument("--gpu", action="store_true")
+
+	args = p.parse_args()
+
+	tok_file = f"{args.tokensdir}{args.text}.tokenizer.json"
+	dat_file = f"{args.tokensdir}{args.text}_bs={args.blocksize}.tar.gz"
+	output_dir = f"{args.ouputdir}{args.text}_bs={args.blocksize}_ns={args.numsteps}"
+
+	if args.gpu:
+		return run_gpu(te=output_dir, tok=tok_file, dat=dat_file, blocksize=args.blocksize, num_steps=args.numsteps)
+	else:
+		return run_cpu(output_dir, tok_file, dat_file, args.blocksize)
+
+if __name__ == '__main__':
+    sys.exit(main())