from aitextgen import aitextgen

if __name__ == '__main__':
    # With your trained model, you can reload the model at any time by
    # providing the folder containing the pytorch_model.bin model weights + the config, and providing the tokenizer.
    ai2 = aitextgen(model_folder="trained_model",
                    tokenizer_file="aitextgen.tokenizer.json")

    ai2.generate(10)