# -*- coding: utf-8 -*-
from setuptools import setup

packages = \
['perceiver',
 'perceiver.data',
 'perceiver.data.text',
 'perceiver.data.vision',
 'perceiver.model',
 'perceiver.model.core',
 'perceiver.model.text',
 'perceiver.model.vision',
 'perceiver.scripts',
 'perceiver.scripts.text',
 'perceiver.scripts.vision']

package_data = \
{'': ['*']}

install_requires = \
['einops>=0.4,<0.5',
 'fairscale>=0.4,<0.5',
 'fsspec[s3]',
 'jsonargparse[signatures]>=4.12,<5.0',
 'pytorch-lightning>=1.7,<2.0',
 'torch-optimizer>=0.3,<0.4',
 'torch>=1.12,<2.0',
 'torchmetrics>=0.9,<0.10']

extras_require = \
{'text': ['datasets>=2.4,<3.0',
          'tokenizers>=0.12,<0.13',
          'transformers>=4.21,<5.0'],
 'vision': ['datasets>=2.4,<3.0', 'torchvision>=0.13,<0.14']}

setup_kwargs = {
    'name': 'perceiver-io',
    'version': '0.7b1',
    'description': 'Perceiver IO',
    'long_description': '# Perceiver, Perceiver IO and Perceiver AR\n\nThis repository is a PyTorch and PyTorch Lightning implementation of\n\n<table>\n  <tr>\n    <td>\n       <b>Perceiver</b>: General Perception with Iterative Attention\n       (<a href="https://arxiv.org/abs/2103.03206">paper</a>,\n        <a href="https://www.youtube.com/watch?v=P_xeshTnPZg">video</a>)\n    </td>\n    <td><img src="docs/images/small-perceiver.png" alt="Perceiver"/></td>\n  </tr>\n  <tr>\n    <td>\n      <b>Perceiver IO</b>: A General Architecture for Structured Inputs & Outputs\n      (<a href="https://arxiv.org/abs/2107.14795">paper</a>,\n       <a href="https://www.deepmind.com/blog/building-architectures-that-can-handle-the-worlds-data">blog post</a>)\n    </td>\n    <td><img src="docs/images/small-perceiver-io.png" alt="Perceiver IO"/></td>\n  </tr>\n  <tr>\n    <td>\n      General-purpose, long-context autoregressive modeling with <b>Perceiver AR</b>\n      (<a href="https://arxiv.org/abs/2202.07765">paper</a>,\n       <a href="https://www.deepmind.com/blog/perceiver-ar-general-purpose-long-context-autoregressive-generation">blog post</a>)\n    </td>\n    <td><img src="docs/images/small-perceiver-ar.png" alt="Perceiver AR"/></td>\n  </tr>\n</table>\n\nModels are written in plain PyTorch and can be trained at scale with [PyTorch Lightning](https://pytorch-lightning.readthedocs.io/en/stable/)\nand the [Lightning CLI](https://pytorch-lightning.readthedocs.io/en/stable/cli/lightning_cli.html). [Pretrained weights](docs/pretrained-models.md)\nof 🤗 Perceiver models are also supported and can be downloaded from the 🤗 Hub. The provided datasets for model training\nare 🤗 [datasets](https://huggingface.co/docs/datasets) wrapped into PyTorch Lightning [data modules](perceiver/data).\nFor NLP tasks, this library supports all 🤗 [fast tokenizers](https://huggingface.co/docs/transformers/fast_tokenizers)\nand the 🤗 Perceiver UTF-8 bytes tokenizer.\n\n## Installation\n\n### Via pip\n\n```shell\npip install perceiver-io[text,vision]\n```\n\n### From sources\n\nInstallation from sources requires a [Miniconda](https://docs.conda.io/en/latest/miniconda.html) and a\n[Poetry](https://python-poetry.org/docs/#installation) (1.2.0 or higher) installation.\n\n```shell\nconda env create -f environment.yml\nconda activate perceiver-io\npoetry install --all-extras\n```\n\n### Docker image\n\n```shell\ndocker pull ghcr.io/krasserm/perceiver-io:latest\n```\n\nSee [Docker image](docs/docker-image.md) for details.\n\n## Documentation\n\n- [Getting started](#getting-started)\n- [Model construction](docs/model-construction.md)\n- [Pretrained models](docs/pretrained-models.md)\n- [Training examples](docs/training-examples.md)\n- [Inference examples](examples/inference.ipynb) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/krasserm/perceiver-io/blob/main/examples/inference.ipynb)\n- [Building blocks](docs/building-blocks.md)\n\n## Getting started\n\nHere\'s a minimal example for autoregressive language modeling with Perceiver AR. A small language model (30.7M parameters)\nis trained on the WikiText-103-raw dataset and then used to generate text from a prompt. Input text is tokenized into\nraw UTF-8 bytes and the model also generates raw UTF-8 bytes.\n\nThe PyTorch model class (`CausalLanguageModel`) and the corresponding PyTorch Lightning wrapper class\n(`LitCausalLanguageModel`) are defined in [perceiver/model/text/clm.py](perceiver/model/text/clm.py) (see also\n[model construction](docs/model-construction.md) for further details). The PyTorch Lightning data module\n(`WikiTextDataModule`) is defined in [perceiver/data/text/wikitext.py](perceiver/data/text/wikitext.py).\n\n### Training\n\n#### Command line\n\nThe script for training a `CausalLanguageModel` on the command line is [perceiver/scripts/text/clm.py](perceiver/scripts/text/clm.py).\nThe constructor signatures of `LitCausalLanguageModel` and `WikiTextDataModule` determine the available `--model.*` and\n`--data.*` command line options. Command line options `--optimizer.*`, `--lr_scheduler.*` and `--trainer.*` configure\nthe optimizer, learning rate scheduler and the PyTorch Lightning [Trainer](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html),\nrespectively.\n\n```shell\npython -m perceiver.scripts.text.clm fit \\\n  --model.num_latents=512 \\\n  --model.num_channels=512 \\\n  --model.num_self_attention_layers=8 \\\n  --model.cross_attention_dropout=0.5 \\\n  --data=WikiTextDataModule \\\n  --data.tokenizer=deepmind/language-perceiver \\\n  --data.add_special_tokens=false \\\n  --data.max_seq_len=4096 \\\n  --data.task=clm \\\n  --data.batch_size=16 \\\n  --optimizer=Adam \\\n  --optimizer.lr=2e-4 \\\n  --lr_scheduler.warmup_steps=200 \\\n  --trainer.accelerator=gpu \\\n  --trainer.devices=1 \\\n  --trainer.max_epochs=5 \\\n  --trainer.accumulate_grad_batches=4\n```\n\nSupported optimizers are those packaged with PyTorch and [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).\nThe `--data.task=clm` option configures the data module to produce data compatible with causal language modeling (other\npossible values are `mlm` for masked language modeling and `clf` for sequence classification). When running this command\nfor the first time, the WikiText dataset is downloaded and preprocessed. To download and preprocess the dataset prior\nto training, run\n\n```shell\npython -m perceiver.scripts.text.preproc wikitext \\\n  --tokenizer=deepmind/language-perceiver \\\n  --add_special_tokens=false \\\n  --max_seq_len=4096 \\\n  --task=clm\n```\n\nwhich is usually faster.\n\n#### Python code\n\nTraining on the command line uses the PyTorch Lightning `Trainer` under the hood. To run the `Trainer` directly from\na Python script, dynamically add a `configure_optimizers` method to `LitCausalLanguageModel`, create instances of\n`LitCausalLanguageModel` and `WikiTextDataModule` and then call `trainer.fit()` with the model and data module as\narguments:\n\n```python\nfrom torch.optim import Adam\n\nfrom perceiver.data.text import WikiTextDataModule, Task\nfrom perceiver.model.text.clm import LitCausalLanguageModel, CausalLanguageModelConfig\nfrom perceiver.scripts.lrs import ConstantWithWarmupLR\n\nimport pytorch_lightning as pl\n\n\ndef configure_optimizers(self):\n    optimizer = Adam(self.parameters(), lr=2e-4)\n    scheduler = ConstantWithWarmupLR(optimizer, warmup_steps=200)\n    return {\n        "optimizer": optimizer,\n        "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1},\n    }\n\n\n# # Add configure_optimizers method to LitCausalLanguageModel (not hard-coded there)\nsetattr(LitCausalLanguageModel, "configure_optimizers", configure_optimizers),\n\n\nif __name__ == \'__main__\':\n    data = WikiTextDataModule(\n        tokenizer="deepmind/language-perceiver",\n        add_special_tokens=False,\n        max_seq_len=4096,\n        task=Task.clm,\n        batch_size=16,\n    )\n\n    config = CausalLanguageModelConfig(\n        vocab_size=data.vocab_size,\n        max_seq_len=data.max_seq_len,\n        num_latents=512,\n        num_channels=512,\n        num_self_attention_layers=8,\n        cross_attention_dropout=0.5,\n    )\n\n    # Create Lightning module of CausalLanguageModel from configuration object\n    lit_model = LitCausalLanguageModel.create(config)\n\n    # Instantiate PyTorch Lightning Trainer\n    trainer = pl.Trainer(accelerator="gpu", devices=1, max_epochs=5, accumulate_grad_batches=4)\n\n    # Train model (will also preprocess dataset if not already done yet)\n    trainer.fit(lit_model, datamodule=data)\n```\n\nThe trained PyTorch model can be accessed with `lit_model.model`. If you prefer to use a custom training loop without\nusing the PyTorch Lightning Trainer, create a plain PyTorch model with `CausalLanguageModel.create(config=...)` and\ntrain it directly as shown in the following simple example:\n\n```python\nfrom perceiver.model.text.clm import CausalLanguageModel\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.optim import Adam\n\ndata = ...\ndata.prepare_data()\ndata.setup()\n\nmodel_config = ...\nmodel = CausalLanguageModel(config=model_config)\nmodel.train()\n\noptim = Adam(model.parameters(), lr=2e-4)\n\n# Simplified training loop compared to previous\n# examples (no gradient accumulation, ...)\nfor epoch in range(5):\n    for labels_ids, input_ids, _ in data.train_dataloader():\n        logits = model(input_ids)\n        loss = F.cross_entropy(logits.permute(0, 2, 1), labels_ids[:, -model_config.num_latents:])\n        loss.backward()\n        optim.step()\n        optim.zero_grad()\n\n# Save trained model\ntorch.save(model.state_dict(), "/path/to/model.pt")\n```\n\n### Inference\n\nFor generating text from a prompt via top-k sampling, `CausalLanguageModel` provides a `generate()` method. The following\nexample first loads a trained model from a checkpoint and then generates text from a short sample prompt. An interactive\ndemo is also available in the [Colab notebook](https://colab.research.google.com/github/krasserm/perceiver-io/blob/main/examples/inference.ipynb).\n\n```python\nfrom perceiver.data.text import TextPreprocessor\nfrom perceiver.model.text.clm import LitCausalLanguageModel\n\n# Load model from a checkpoint that has been written by the PyTorch Lightning Trainer\nmodel = LitCausalLanguageModel.load_from_checkpoint("/path/to/checkpoint").model.eval()\n\n# Alternatively, load the model\'s state_dict directly\n#model = CausalLanguageModel(config=model_config).eval()\n#model.load_state_dict(torch.load("/path/to/model.pt"))\n\n# Create a text preprocessor  \npreproc = TextPreprocessor(tokenizer="deepmind/language-perceiver", max_seq_len=4096, add_special_tokens=False)\n\n# Convert text to model input\nprompt, _ = preproc.preprocess("A man was reading a book on a sunny day until he sudden")\n\n# Continue prompt via top-k sampling where k = f(vocab_size, threshold)\ngenerated = model.generate(num=512, prompt=prompt[None, ...], threshold=0.9)\n\n# Decode model output using preprocessor\'s tokenizer\ngenerated_text = preproc.tokenizer.decode(generated[0])\n```\n\n## Other implementations\n\n- [Perceiver](https://paperswithcode.com/paper/perceiver-general-perception-with-iterative#code)\n- [Perceiver IO](https://paperswithcode.com/paper/perceiver-io-a-general-architecture-for#code)\n- [Perceiver AR](https://paperswithcode.com/paper/general-purpose-long-context-autoregressive#code)\n',
    'author': 'Martin Krasser',
    'author_email': 'krasserm@googlemail.com',
    'maintainer': 'None',
    'maintainer_email': 'None',
    'url': 'https://github.com/krasserm/perceiver-io',
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'extras_require': extras_require,
    'python_requires': '>=3.7,<4.0',
}


setup(**setup_kwargs)
