{"id":15601037,"url":"https://github.com/lucidrains/electra-pytorch","last_synced_at":"2025-04-06T01:07:14.216Z","repository":{"id":50309700,"uuid":"285105911","full_name":"lucidrains/electra-pytorch","owner":"lucidrains","description":"A simple and working implementation of Electra, the fastest way to pretrain language models from scratch, in Pytorch","archived":false,"fork":false,"pushed_at":"2023-06-12T21:28:57.000Z","size":95,"stargazers_count":225,"open_issues_count":10,"forks_count":46,"subscribers_count":9,"default_branch":"master","last_synced_at":"2025-03-30T00:07:17.673Z","etag":null,"topics":["artificial-intelligence","deep-learning","pre-training","transformer"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2020-08-04T21:38:03.000Z","updated_at":"2025-03-27T19:42:41.000Z","dependencies_parsed_at":"2024-10-23T01:32:44.108Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/electra-pytorch","commit_stats":{"total_commits":59,"total_committers":2,"mean_commits":29.5,"dds":"0.13559322033898302","last_synced_commit":"5b8bae5c3575b7529891c1b878c24688f57d7ca1"},"previous_names":[],"tags_count":4,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Felectra-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Felectra-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Felectra-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Felectra-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/electra-pytorch/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247419860,"owners_count":20936012,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","deep-learning","pre-training","transformer"],"created_at":"2024-10-03T02:12:44.419Z","updated_at":"2025-04-06T01:07:14.193Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"\u003cimg src=\"./electra.png\"\u003e\u003c/img\u003e\n\n## Electra - Pytorch\n\nA simple working wrapper for fast pretraining of language models as detailed in \u003ca href=\"https://arxiv.org/abs/2003.10555\"\u003ethis paper\u003c/a\u003e. It speeds up training (in comparison to normal masked language modeling) by a factor of 4x, and eventually reaches better performance if trained for even longer. Special thanks to \u003ca href=\"https://github.com/enijkamp\"\u003eErik Nijkamp\u003c/a\u003e for taking the time to replicate the results for GLUE.\n\n## Install\n\n```bash\n$ pip install electra-pytorch\n```\n\n## Usage\n\nThe following example uses `reformer-pytorch`, which is available to be pip installed.\n\n```python\nimport torch\nfrom torch import nn\nfrom reformer_pytorch import ReformerLM\n\nfrom electra_pytorch import Electra\n\n# (1) instantiate the generator and discriminator, making sure that the generator is roughly a quarter to a half of the size of the discriminator\n\ngenerator = ReformerLM(\n    num_tokens = 20000,\n    emb_dim = 128,\n    dim = 256,              # smaller hidden dimension\n    heads = 4,              # less heads\n    ff_mult = 2,            # smaller feed forward intermediate dimension\n    dim_head = 64,\n    depth = 12,\n    max_seq_len = 1024\n)\n\ndiscriminator = ReformerLM(\n    num_tokens = 20000,\n    emb_dim = 128,\n    dim = 1024,\n    dim_head = 64,\n    heads = 16,\n    depth = 12,\n    ff_mult = 4,\n    max_seq_len = 1024\n)\n\n# (2) weight tie the token and positional embeddings of generator and discriminator\n\ngenerator.token_emb = discriminator.token_emb\ngenerator.pos_emb = discriminator.pos_emb\n# weight tie any other embeddings if available, token type embeddings, etc.\n\n# (3) instantiate electra\n\ntrainer = Electra(\n    generator,\n    discriminator,\n    discr_dim = 1024,           # the embedding dimension of the discriminator\n    discr_layer = 'reformer',   # the layer name in the discriminator, whose output would be used for predicting token is still the same or replaced\n    mask_token_id = 2,          # the token id reserved for masking\n    pad_token_id = 0,           # the token id for padding\n    mask_prob = 0.15,           # masking probability for masked language modeling\n    mask_ignore_token_ids = []  # ids of tokens to ignore for mask modeling ex. (cls, sep)\n)\n\n# (4) train\n\ndata = torch.randint(0, 20000, (1, 1024))\n\nresults = trainer(data)\nresults.loss.backward()\n\n# after much training, the discriminator should have improved\n\ntorch.save(discriminator, f'./pretrained-model.pt')\n```\n\nIf you would rather not have the framework auto-magically intercept the hidden output of the discriminator, you can pass in the discriminator (with the extra linear [dim x 1]) by yourself with the following.\n\n```python\nimport torch\nfrom torch import nn\nfrom reformer_pytorch import ReformerLM\n\nfrom electra_pytorch import Electra\n\n# (1) instantiate the generator and discriminator, making sure that the generator is roughly a quarter to a half of the size of the discriminator\n\ngenerator = ReformerLM(\n    num_tokens = 20000,\n    emb_dim = 128,\n    dim = 256,              # smaller hidden dimension\n    heads = 4,              # less heads\n    ff_mult = 2,            # smaller feed forward intermediate dimension\n    dim_head = 64,\n    depth = 12,\n    max_seq_len = 1024\n)\n\ndiscriminator = ReformerLM(\n    num_tokens = 20000,\n    emb_dim = 128,\n    dim = 1024,\n    dim_head = 64,\n    heads = 16,\n    depth = 12,\n    ff_mult = 4,\n    max_seq_len = 1024,\n    return_embeddings = True\n)\n\n# (2) weight tie the token and positional embeddings of generator and discriminator\n\ngenerator.token_emb = discriminator.token_emb\ngenerator.pos_emb = discriminator.pos_emb\n# weight tie any other embeddings if available, token type embeddings, etc.\n\n# (3) instantiate electra\n\ndiscriminator_with_adapter = nn.Sequential(discriminator, nn.Linear(1024, 1))\n\ntrainer = Electra(\n    generator,\n    discriminator_with_adapter,\n    mask_token_id = 2,          # the token id reserved for masking\n    pad_token_id = 0,           # the token id for padding\n    mask_prob = 0.15,           # masking probability for masked language modeling\n    mask_ignore_token_ids = []  # ids of tokens to ignore for mask modeling ex. (cls, sep)\n)\n\n# (4) train\n\ndata = torch.randint(0, 20000, (1, 1024))\n\nresults = trainer(data)\nresults.loss.backward()\n\n# after much training, the discriminator should have improved\n\ntorch.save(discriminator, f'./pretrained-model.pt')\n```\n\n## Important details for successful training\n\nThe generator should be roughly a quarter to at most one half of the discriminator's size for effective training. Any greater and the generator will be too good and the adversarial game collapses. This was done by reducing the hidden dimension, feed forward hidden dimension, and number of attention heads in the paper.\n\n## Testing\n\n```bash\n$ python setup.py test\n```\n\n## Training\n\n1. Download the [OpenWebText](https://github.com/jcpeterson/openwebtext) dataset.\n\n```bash\n$ mkdir data\n$ cd data\n$ pip3 install gdown\n$ gdown --id 1EA5V0oetDCOke7afsktL_JDQ-ETtNOvx\n$ tar -xf openwebtext.tar.xz\n$ wget https://storage.googleapis.com/electra-data/vocab.txt\n$ cd ..\n```\n\n2. Tokenize dataset.\n\n```bash\n$ python pretraining/openwebtext/preprocess.py\n```\n\n3. Pre-train.\n\n```bash\n$ python pretraining/openwebtext/pretrain.py\n```\n\n4. Download GLUE dataset.\n\n```bash\n$ python examples/glue/download.py \n```\n\n5. Fine-tune on the MRPC sub-task of the GLUE benchmark.\n\n```bash\n$ python examples/glue/run.py --model_name_or_path output/yyyy-mm-dd-hh-mm-ss/ckpt/200000\n```\n\n## Citations\n\n```bibtex\n@misc{clark2020electra,\n    title={ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators},\n    author={Kevin Clark and Minh-Thang Luong and Quoc V. Le and Christopher D. Manning},\n    year={2020},\n    eprint={2003.10555},\n    archivePrefix={arXiv},\n    primaryClass={cs.CL}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Felectra-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Felectra-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Felectra-pytorch/lists"}