Package: aitextgen
aitextgenA robust Python tool for text-based AI training and generation usingOpenAI'sGPT-2andEleutherAI'sGPT Neo/GPT-3architecture.aitextgen is a Python package that leveragesPyTorch,Hugging Face Transformersandpytorch-lightningwith specific optimizations for text generation using GPT-2, plusmanyadded features. It is the successor totextgenrnnandgpt-2-simple, taking the best of both packages:Finetunes on a pretrained 124M/355M/774M GPT-2 model from OpenAI or a 125M/350M GPT Neo model from EleutherAI...or create your own GPT-2/GPT Neo model + tokenizer and train from scratch!Generates text faster than gpt-2-simple and with better memory efficiency!With Transformers, aitextgen preserves compatibility with the base package, allowing you to use the model for other NLP tasks, download custom GPT-2 models from the HuggingFace model repository, and upload your own models! Also, it uses the includedgenerate()function to allow a massive amount of control over the generated text.With pytorch-lightning, aitextgen trains models not just on CPUs and GPUs, but alsomultipleGPUs and (eventually) TPUs! It also includes a pretty training progress bar, with the ability to add optional loggers.The input dataset is its own object, allowing you to not only easily encode megabytes of data in seconds, cache, and compress it on a local computer before transporting to a remote server, but you are able tomergedatasets without biasing the resulting dataset, orcross-trainon multiple datasets to create blended output.You can read more about aitextgenin the documentation!DemoYou can play with aitextgenfor freewith powerful GPUs using these Colaboratory Notebooks!Finetune OpenAI's 124M GPT-2 model (or GPT Neo) on your own dataset (GPU)Train a GPT-2 model + tokenizer from scratch (GPU)You can also play with customRedditandHacker Newsdemo models on your own PC.Installationaitextgen can be installedfrom PyPI:pip3installaitextgenQuick ExamplesHere's how you can quickly test out aitextgen on your own computer, even if you don't have a GPU!For generating text from a pretrained GPT-2 model:fromaitextgenimportaitextgen# Without any parameters, aitextgen() will download, cache, and load the 124M GPT-2 "small" modelai=aitextgen()ai.generate()ai.generate(n=3,max_length=100)ai.generate(n=3,prompt="I believe in unicorns because",max_length=100)ai.generate_to_file(n=10,prompt="I believe in unicorns because",max_length=100,temperature=1.2)You can also generate from the command line:aitextgengenerate
aitextgengenerate--prompt"I believe in unicorns because"--to_fileFalseWant to train your own mini GPT-2 model on your own computer? You can follow alongin this Jupyter Notebookor, download thistext file of Shakespeare's plays, cd to that directory in a Terminal, open up apython3console and go:fromaitextgen.TokenDatasetimportTokenDatasetfromaitextgen.tokenizersimporttrain_tokenizerfromaitextgen.utilsimportGPT2ConfigCPUfromaitextgenimportaitextgen# The name of the downloaded Shakespeare text for trainingfile_name="input.txt"# Train a custom BPE Tokenizer on the downloaded text# This will save one file: `aitextgen.tokenizer.json`, which contains the# information needed to rebuild the tokenizer.train_tokenizer(file_name)tokenizer_file="aitextgen.tokenizer.json"# GPT2ConfigCPU is a mini variant of GPT-2 optimized for CPU-training# e.g. the # of input tokens here is 64 vs. 1024 for base GPT-2.config=GPT2ConfigCPU()# Instantiate aitextgen using the created tokenizer and configai=aitextgen(tokenizer_file=tokenizer_file,config=config)# You can build datasets for training by creating TokenDatasets,# which automatically processes the dataset with the appropriate size.data=TokenDataset(file_name,tokenizer_file=tokenizer_file,block_size=64)# Train the model! It will save pytorch_model.bin periodically and after completion to the `trained_model` folder.# On a 2020 8-core iMac, this took ~25 minutes to run.ai.train(data,batch_size=8,num_steps=50000,generate_every=5000,save_every=5000)# Generate text from it!ai.generate(10,prompt="ROMEO:")# 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,prompt="ROMEO:")Want to run aitextgen and finetune GPT-2? Use the Colab notebooks in the Demos section, orfollow the documentationto get more information and learn some helpful tips!Known IssuesTPUs cannot be used to train a model: although youcantrain an aitextgen model on TPUs by settingn_tpu_cores=8in an appropriate runtime, and the training loss indeed does decrease, there are a number of miscellaneous blocking problems. [Tracking GitHub Issue]Upcoming FeaturesThe current release (v0.5.X) of aitextgenis considered to be a beta, targeting the most common use cases. The Notebooks and examples written so far are tested to work, but more fleshing out of the docs/use cases will be done over the next few months in addition to fixing the known issues noted above.The next versions of aitextgen (and one of the reasons I made this package in the first place) will have native support forschema-based generation. (Seethis repofor a rough proof-of-concept.)Additionally, I plan to develop an aitextgenSaaSto allow anyone to run aitextgen in the cloud and build APIs/Twitter+Slack+Discord bots with just a few clicks. (The primary constraint is compute cost; if any venture capitalists are interested in funding the development of such a service, let me know.)I've listed more tentative features in theUPCOMINGdocument.Ethicsaitextgen is a tool primarily intended to help facilitate creative content. It is not a tool intended to deceive. Although parody accounts are an obvious use case for this package, make sure you areas upfront as possiblewith the methodology of the text you create. This includes:State that the text was generated using aitextgen and/or a GPT-2 model architecture. (A link to this repo would be a bonus!)If parodying a person, explicitly state that it is a parody, and reference who it is parodying.If the generated text is human-curated, or if it's unsupervised random output.Indicating who is maintaining/curating the AI-generated text.Make a good-faith effort to remove overfit output from the generated text that matches the input text verbatim.It's fun to anthropomorphise the nameless "AI" as an abstract genius, but part of the reason I made aitextgen (and all my previous text-generation projects) is to make the technology more accessible and accurately demonstrate both its promise, and its limitations.Any AI text generation projects that are deliberately deceptive may be disavowed.Maintainer/CreatorMax Woolf (@minimaxir)Max's open-source projects are supported by hisPatreonandGitHub Sponsors. If you found this project helpful, any monetary contributions to the Patreon are appreciated and will be put to good creative use.LicenseMIT