[WIP] Add GeoV model

docs/source/en/model_doc/geov.mdx

@@ -0,0 +1,68 @@
+<!--Copyright 2023 The Better Planet Investments, labml.ai and The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# GeoV
+
+## Overview
+
+The GeoV model was designed by Georges Harik and uses [Rotary Positional Embeddings with Relative distances (RoPER)](http://research.labml.ai/RoPER.html) by [Georges Hark](https://twitter.com/ghark) and [Varuna Jayasiri](https://twitter.com/vpj).
+
+[RoPER](http://research.labml.ai/RoPER.html), in addition to using relative positions in the attention score calculation by RoPE embeddings, adds relative positional information explicitly to value embeddings. Specifically, it incorporates the relative positions of the tokens paid attention to. RoPER has given better performance in some algorithmic tasks, and seems comparable to RoPE in language modeling.
+
+The GeoV tokenizer uses [SentencePiece](https://github.com/google/sentencepiece) [unigram language model](https://arxiv.org/abs/1804.10959) and tokenizes symbols, digits and new line characters separately, in order to achieve better performance on mathematical content and code.
+
+This model was contributed by [gharik](https://huggingface.co/gharik) and [vpj](https://huggingface.co/vpj).
+
+We have shared 9B parameter pre-trained model at [GeoV/GeoV-9b](https://huggingface.co/GeoV/GeoV-9b).
+We plan to release checkpoints around every 20b tokens trained from here until around 300b tokens.
+We will also train smaller and larger versions. Our aim is to have broadly available smaller and larger models.
+
+## Generation
+
+The `generate()` method can be used to generate text using GeoV model.
+
+```python
+>>> from transformers import GeoVForCausalLM, GeoVTokenizer
+
+>>> model = GeoVForCausalLM.from_pretrained("GeoV/GeoV-9b")
+>>> tokenizer = GeoVTokenizer.from_pretrained("GeoV/GeoV-9b")
+
+>>> prompt = "In mathematics, topology is the study of"
+
+>>> input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+
+>>> gen_tokens = model.generate(
+...     input_ids,
+...     do_sample=True,
+...     temperature=0.9,
+...     max_length=100,
+... )
+>>> gen_text = tokenizer.batch_decode(gen_tokens)[0]
+```
+
+## GeoVConfig
+
+[[autodoc]] GeoVConfig
+
+## GeoVTokenizer
+
+[[autodoc]] GeoVTokenizer
+
+## GeoVModel
+
+[[autodoc]] GeoVModel
+    - forward
+
+## GeoVForCausalLM
+
+[[autodoc]] GeoVForCausalLM
+    - forward

src/transformers/__init__.py

@@ -300,6 +300,7 @@
     "models.gpt2": ["GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPT2Config", "GPT2Tokenizer"],
     "models.gpt_neo": ["GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTNeoConfig"],
     "models.gpt_neox": ["GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTNeoXConfig"],
+    "models.geov": ["GEOV_PRETRAINED_CONFIG_ARCHIVE_MAP", "GeoVConfig"],
     "models.gpt_neox_japanese": ["GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTNeoXJapaneseConfig"],
     "models.gpt_sw3": [],
     "models.gptj": ["GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTJConfig"],
@@ -732,6 +733,7 @@
     _import_structure["models.funnel"].append("FunnelTokenizerFast")
     _import_structure["models.gpt2"].append("GPT2TokenizerFast")
     _import_structure["models.gpt_neox"].append("GPTNeoXTokenizerFast")
+    _import_structure["models.geov"].append("GeoVTokenizer")
     _import_structure["models.gpt_neox_japanese"].append("GPTNeoXJapaneseTokenizer")
     _import_structure["models.herbert"].append("HerbertTokenizerFast")
     _import_structure["models.layoutlm"].append("LayoutLMTokenizerFast")
@@ -1657,6 +1659,15 @@
             "GPTNeoXPreTrainedModel",
         ]
     )
+    _import_structure["models.geov"].extend(
+        [
+            "GEOV_PRETRAINED_MODEL_ARCHIVE_LIST",
+            "GeoVForCausalLM",
+            "GeoVLayer",
+            "GeoVModel",
+            "GeoVPreTrainedModel",
+        ]
+    )
     _import_structure["models.gpt_neox_japanese"].extend(
         [
             "GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST",
@@ -3956,6 +3967,7 @@
     from .models.gpt2 import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP, GPT2Config, GPT2Tokenizer
     from .models.gpt_neo import GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoConfig
     from .models.gpt_neox import GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoXConfig
+    from .models.geov import GEOV_PRETRAINED_CONFIG_ARCHIVE_MAP, GeoVConfig
     from .models.gpt_neox_japanese import GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoXJapaneseConfig
     from .models.gptj import GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTJConfig
     from .models.gptsan_japanese import (
@@ -4353,6 +4365,7 @@
         from .models.funnel import FunnelTokenizerFast
         from .models.gpt2 import GPT2TokenizerFast
         from .models.gpt_neox import GPTNeoXTokenizerFast
+        from .models.geov import GeoVTokenizer
         from .models.gpt_neox_japanese import GPTNeoXJapaneseTokenizer
         from .models.herbert import HerbertTokenizerFast
         from .models.layoutlm import LayoutLMTokenizerFast
@@ -5117,6 +5130,13 @@
             GPTNeoXModel,
             GPTNeoXPreTrainedModel,
         )
+        from .models.geov import (
+            GEOV_PRETRAINED_MODEL_ARCHIVE_LIST,
+            GeoVForCausalLM,
+            GeoVLayer,
+            GeoVModel,
+            GeoVPreTrainedModel,
+        )
         from .models.gpt_neox_japanese import (
             GPT_NEOX_JAPANESE_PRETRAINED_MODEL_ARCHIVE_LIST,
             GPTNeoXJapaneseForCausalLM,

src/transformers/models/__init__.py

@@ -84,6 +84,7 @@
     gpt2,
     gpt_neo,
     gpt_neox,
+    geov,
     gpt_neox_japanese,
     gpt_sw3,
     gptj,

src/transformers/models/auto/configuration_auto.py

@@ -93,6 +93,7 @@
         ("gpt2", "GPT2Config"),
         ("gpt_neo", "GPTNeoConfig"),
         ("gpt_neox", "GPTNeoXConfig"),
+        ("geov", "GeoVConfig"),
         ("gpt_neox_japanese", "GPTNeoXJapaneseConfig"),
         ("gptj", "GPTJConfig"),
         ("gptsan-japanese", "GPTSanJapaneseConfig"),
@@ -273,6 +274,7 @@
         ("gpt2", "GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
         ("gpt_neo", "GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
         ("gpt_neox", "GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP"),
+        ("geov", "GEOV_PRETRAINED_CONFIG_ARCHIVE_MAP"),
         ("gpt_neox_japanese", "GPT_NEOX_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
         ("gptj", "GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP"),
         ("gptsan-japanese", "GPTSAN_JAPANESE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
@@ -454,6 +456,7 @@
         ("gpt2", "OpenAI GPT-2"),
         ("gpt_neo", "GPT Neo"),
         ("gpt_neox", "GPT NeoX"),
+        ("geov", "GeoV"),
         ("gpt_neox_japanese", "GPT NeoX Japanese"),
         ("gptj", "GPT-J"),
         ("gptsan-japanese", "GPTSAN-japanese"),

src/transformers/models/auto/modeling_auto.py

@@ -91,6 +91,7 @@
         ("gpt2", "GPT2Model"),
         ("gpt_neo", "GPTNeoModel"),
         ("gpt_neox", "GPTNeoXModel"),
+        ("geov", "GeoVModel"),
         ("gpt_neox_japanese", "GPTNeoXJapaneseModel"),
         ("gptj", "GPTJModel"),
         ("gptsan-japanese", "GPTSanJapaneseForConditionalGeneration"),
@@ -293,6 +294,7 @@
         ("gpt2", "GPT2LMHeadModel"),
         ("gpt_neo", "GPTNeoForCausalLM"),
         ("gpt_neox", "GPTNeoXForCausalLM"),
+        ("geov", "GeoVForCausalLM"),
         ("gpt_neox_japanese", "GPTNeoXJapaneseForCausalLM"),
         ("gptj", "GPTJForCausalLM"),
         ("gptsan-japanese", "GPTSanJapaneseForConditionalGeneration"),
@@ -361,6 +363,7 @@
         ("gpt2", "GPT2LMHeadModel"),
         ("gpt_neo", "GPTNeoForCausalLM"),
         ("gpt_neox", "GPTNeoXForCausalLM"),
+        ("geov", "GeoVForCausalLM"),
         ("gpt_neox_japanese", "GPTNeoXJapaneseForCausalLM"),
         ("gptj", "GPTJForCausalLM"),
         ("llama", "LlamaForCausalLM"),

src/transformers/models/auto/tokenization_auto.py

@@ -153,6 +153,7 @@
             ("gpt2", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
             ("gpt_neo", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
             ("gpt_neox", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+            ("geov", (None, "GeoVTokenizer" if is_tokenizers_available() else None)),
             ("gpt_neox_japanese", ("GPTNeoXJapaneseTokenizer", None)),
             ("gptj", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
             ("gptsan-japanese", ("GPTSanJapaneseTokenizer", None)),

src/transformers/models/geov/__init__.py

@@ -0,0 +1,74 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...file_utils import _LazyModule, is_tokenizers_available, is_torch_available
+from ...utils import OptionalDependencyNotAvailable
+
+
+_import_structure = {"configuration_geov": ["GEOV_PRETRAINED_CONFIG_ARCHIVE_MAP", "GeoVConfig"]}
+
+try:
+    if not is_tokenizers_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["tokenization_geov"] = ["GeoVTokenizer"]
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_geov"] = [
+        "GEOV_PRETRAINED_MODEL_ARCHIVE_LIST",
+        "GeoVForCausalLM",
+        "GeoVLayer",
+        "GeoVModel",
+        "GeoVPreTrainedModel",
+    ]
+
+
+if TYPE_CHECKING:
+    from .configuration_geov import GEOV_PRETRAINED_CONFIG_ARCHIVE_MAP, GeoVConfig
+
+    try:
+        if not is_tokenizers_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .tokenization_geov import GeoVTokenizer
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_geov import (
+            GEOV_PRETRAINED_MODEL_ARCHIVE_LIST,
+            GeoVForCausalLM,
+            GeoVLayer,
+            GeoVModel,
+            GeoVPreTrainedModel,
+        )
+
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

src/transformers/models/geov/configuration_geov.py

@@ -0,0 +1,108 @@
+# coding=utf-8
+# Copyright 2023 Better Planet Investments and labml.ai team. ALl rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" GeoV model configuration"""
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+GEOV_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "GeoV/GeoV-9b": "https://huggingface.co/GeoV/GeoV-9b/resolve/main/config.json",
+}
+
+
+class GeoVConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`GeoVModel`]. It is used to instantiate an
+    GeoV model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the GeoV
+    [GeoV/GeoV-9b](https://huggingface.co/GeoV/GeoV-9b) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 50432):
+            Vocabulary size of the GeoV model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`GeoVModel`].
+        hidden_size (`int`, *optional*, defaults to 6144):
+            Dimension of the encoder layers and the pooler layer.
+        num_hidden_layers (`int`, *optional*, defaults to 44):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 64):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (`int`, *optional*, defaults to 24576):
+            Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        rotary_emb_base (`int`, *optional*, defaults to 10000)
+            base for computing rotary embeddings frequency
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        layer_norm_eps (`float`, *optional*, defaults to 1e-4):
+            The epsilon used by the layer normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        use_extra_biases_ffn (`bool`, *optional*, defaults to `False`):
+            Whether or not to have extra bias parameters in the final layer of FFN modules.
+        Example:
+
+    ```python
+    >>> from transformers import GeoVConfig, GeoVModel
+
+    >>> # Initializing a GeoV configuration
+    >>> configuration = GeoVConfig()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = GeoVModel(configuration)  # doctest: +SKIP
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config  # doctest: +SKIP
+    ```"""
+    model_type = "geov"
+
+    def __init__(
+        self,
+        vocab_size=65536,
+        hidden_size=1024 * 5,
+        num_hidden_layers=32,
+        num_attention_heads=40,
+        intermediate_size=1024 * 5 * 4,
+        layer_norm_eps=1e-4,
+        rotary_emb_base=10000,
+        max_position_embeddings=2049,
+        use_extra_biases_ffn=False,
+        use_cache=True,
+        bos_token_id=0,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        super().__init__(
+            bos_token_id=bos_token_id, eos_token_id=eos_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
+        )
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.rotary_emb_base = rotary_emb_base
+        self.use_cache = use_cache
+        self.layer_norm_eps = layer_norm_eps
+        self.use_extra_biases_ffn = use_extra_biases_ffn