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

packages = \
['gato']

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

install_requires = \
['einops', 'torch', 'zetascale']

setup_kwargs = {
    'name': 'gato-torch',
    'version': '0.0.1',
    'description': 'Gato: A Generalist Agent',
    'long_description': '[![Multi-Modality](agorabanner.png)](https://discord.gg/qUtxnK2NMf)\n\n<h1 align="center">Gato: A Generalist Agent</h1>\n\n[[Deepmind Publication]](https://www.deepmind.com/publications/a-generalist-agent)\n[[arXiv Paper]](https://arxiv.org/pdf/2205.06175.pdf)\n\nThis repository contains Deepmind\'s Gato architecture imitation in torch.\n\nSince Deepmind only mentions parts of the architecture in its paper, We still don\'t know much about the model.<br>\nHowever, I believe the paper is enough to imitate the architecture, I\'m trying to do that with the open source community\'s help.\n\nCurrently, the repository supports the following operations:\n- Gato (via [`Gato`](/gato/gato/models/__init__.py))\n- Transformer (via [`Transformer`](/gato/gato/models/transformer.py))\n- Patch Position Encodings (via [`PatchPositionEncoding`](/gato/gato/models/embedding.py))\n- Embedding Function (via [`ResidualEmbedding`](gato/gato/models/embedding.py))\n- Local Observation Position Encodings (via [`LocalPositionEncoding`](gato/gato/models/embedding.py))\n- Tokenizing Continuous Values (via [`ContinuousValueTokenizer`](gato/gato/models/tokenizers.py))\n- Shared Embedding (via [`DiscreteEmbedding`](gato/gato/models/embedding.py))\n\nAction tokens are still a mystery in the paper, I need your help.\n\nHowever, the repository lacks the following miscellaneous.\n- Datasets (most important, Issue: [#1](/datasets/README.md)\n- <s>Pre-trained tokenizers</s> (No longer required because of E2E model)\n- Training strategy (E2E, WIP)\n\nBut, you can still explore the basic architecture of the Gato based on the paper.\n\n### Usage\nThere are 2 methods, git clone + pip:\n\n##### Method1\n\n`git clone https://github.com/kyegomez/GATO.git`\n\n`cd GATO`\n\n`pip install -r requirements.txt`\n\nCreate new file:\n\n\n\n#### Method2\n```bash\n$ pip install gato\n```\n```python\nimport torch\nfrom gato.gato import Gato, GatoConfig\n\n#create model instance\nconfig = GatoConfig.small()\ngato = Gato(config)\n\n\n#fake inputs for Gato\ninput_dim = config.input_dim\ninput_ids = torch.cat([\n    torch.rand((1, 1, input_dim)) for _ in range(20)] + # 20 image patches\n    [torch.full((1, 1, input_dim), 0.25), #continous value]\n     torch.full((1, 1, input_dim), 624.0)] + #discrete (actions, texts)\n     [torch.rand((1, 1, input_dim)) for _ in range(20)] + #20 image patches\n     [torch.full((1, 1, input_dim), 0.12), #continous value\n      torch.full((1, 1, input_dim), 295.0)], #discrete( actions, text)\n      dim=1)\n\nencoding = torch.tensor([\n    [0, 0, 0, 0, 1, 2, 0, 0, 0, 0, 1, 2]\n])\n\nrow_pos = (\n    torch.tensor([[0.00, 0.25, 0.50, 0.75, 0, 0, 0.00, 0.25, 0.50, 0.75, 0, 0]]),  # pos_from\n    torch.tensor([[0.25, 0.50, 0.75, 1.00, 0, 0, 0.25, 0.50, 0.75, 1.00, 0, 0]])  # pos_to\n)\n\ncol_pos = (\n    torch.tensor([[0.00, 0.00, 0.00, 0.80, 0, 0, 0.00, 0.00, 0.00, 0.80, 0, 0]]),  # pos_from\n    torch.tensor([[0.20, 0.20, 0.20, 1.00, 0, 0, 0.20, 0.20, 0.20, 1.00, 0, 0]])  # pos_to\n)\n\n\nobs = (\n    torch.tensor([[ 0,  1,  2, 19, 20, 21,  0,  1,  2, 19, 20, 21]]),  # obs token\n    torch.tensor([[ 1,  1,  1,  1,  1,  0,  1,  1,  1,  1,  1,  0]])  # obs token masking (for action tokens)\n)\n\n\nhidden_states = gato((input_ids, (encoding, row_pos, col_pos), obs))\n```\n\n\n\n### Dataset and Model Architecture\n<picture>\n  <source media="(prefers-color-scheme: dark)" srcset="https://user-images.githubusercontent.com/5837620/215323793-7f7bcfdb-d8be-40d3-8e58-a053511f95d5.png">\n  <img alt="gato dataset and model architecture" src="https://user-images.githubusercontent.com/5837620/215323795-3a433516-f5ca-4272-9999-3df87ae521ba.png">\n</picture>\n\n## Paper Reviews\n\n### Full Episode Sequence\n\n<picture>\n    <source media="(prefers-color-scheme: dark)" srcset="https://user-images.githubusercontent.com/5837620/175756389-31d183c9-054e-4829-93a6-df79781ca212.png">\n    <img alt="gato dataset architecture" src="https://user-images.githubusercontent.com/5837620/175756409-75605dbc-7756-4509-ba93-c0ad08eea309.png">\n</picture>\n\n### Architecture Variants\n\n> Appendix C.1. Transformer Hyperparameters\n\nIn the paper, Deepmind tested Gato with 3 architecture variants, `1.18B`, `364M`, and `79M`.<br>\nI have named them as `large()`, `baseline()` and `small()` respectively in `GatoConfig`.\n\n| Hyperparameters          | Large(1.18B) | Baseline(364M) | Small(79M) |\n|--------------------------|--------------|----------------|------------|\n| Transformer blocks       | 24           | 12             | 8          |\n| Attention heads          | 16           | 12             | 24         |\n| Layer width              | 2048         | 1536           | 768        |\n| Feedforward hidden size  | 8192         | 6144           | 3072       |\n| Key/value size           | 128          | 128            | 32         |\n\n\n### Residual Embedding\n\n> Appendix C.2. Embedding Function\n\nThere are no mentions that how many residual networks must be stacked for token embeddings.<br>\nTherefore, I remain configurable in `GatoConfig`.\n\nWhatever how many residual layers are existing, full-preactivation is a key.\n\nThe blocks are consisted of:\n- Version 2 ResNet architecture (based on ResNet50V2)\n- GroupNorm (instead of LayerNorm)\n- GeLU (instead of ReLU)\n\n### Position Encodings\n\n> Appendix C.3. Position Encodings\n\n#### Patch Position Encodings\n\nLike [Vision Transformer (ViT)](https://github.com/google-research/vision_transformer) by Google, Gato takes the input images as raster-ordered 16x16 patches.<br>\nUnlike the Vision Transformer model, however, Gato divides its patch encoding strategy into 2 phases, training and evaluation.\n\nFor high-performance computation in TensorFlow, I have used the following expressions.\n\n$C$ and $R$ mean column and row-wise, and $F$ and $T$ mean `from` and `to` respectively.\n\n$$\n\\begin{align}\n  v^R_F &= \\begin{bmatrix}\n    0 & 32 & 64 & 96\n  \\end{bmatrix} \\\\\n  v^R_T &= \\begin{bmatrix}\n    32 & 64 & 96 & 128\n  \\end{bmatrix} \\\\\n  v^C_F &= \\begin{bmatrix}\n    0 & 26 & 51 & 77 & 102\n  \\end{bmatrix} \\\\\n  v^C_T &= \\begin{bmatrix}\n    26 & 51 & 77 & 102 & 128\n  \\end{bmatrix} \\\\\n  \\\\\n  P_R &= \\begin{cases}\n    \\mathsf{if} \\ \\mathsf{training} & v^R_F + \\mathsf{uniform}(v^R_T - v^R_F) \\\\\n    \\mathsf{otherwise} & \\mathsf{round}(\\frac{v^R_F + v^R_T}{2})\n  \\end{cases} \\\\\n  P_C &= \\begin{cases}\n    \\mathsf{if} \\ \\mathsf{training} & v^C_F + \\mathsf{uniform}(v^C_T - v^C_F) \\\\\n    \\mathsf{otherwise} & \\mathsf{round}(\\frac{v^C_F + v^C_T}{2})\n  \\end{cases} \\\\\n  \\\\\n  E^R_P &= P_R \\cdot 1^{\\mathsf{T}}_C \\\\\n  E^C_P &= 1^{\\mathsf{T}}_R \\cdot P_C \\\\\n  \\\\\n  \\therefore E &= E_I + E^R_P + E^C_P\n\\end{align}\n$$\n\n#### Local Observation Position Encodings\n\nIn the definition of Appendix B., text tokens, image patch tokens, and discrete & continuous values are observation tokens<br>\nWhen Gato receives those values, they must be encoded with their own (local) time steps.\n\n\n## Contributing\n[We welcome all contributions, please either submit a pull request or submit issues in the Agora discord](https://discord.gg/qUtxnK2NMf)\n\n## License\nLicensed under the [MIT license](/LICENSE).\n\n# Roadmap:\n\n* Get functional prototype\n\n* Integrate ALIBI, multi query, qk norm and other SOTA stuff\n\n* integrate action tokens\n\n',
    'author': 'Kye Gomez',
    'author_email': 'kye@apac.ai',
    'maintainer': 'None',
    'maintainer_email': 'None',
    'url': 'https://github.com/kyegomez/GATO',
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.10,<4.0',
}


setup(**setup_kwargs)
