Add h2oGPT based on Falcon 40B

config.json

@@ -0,0 +1,39 @@
+{
+  "_name_or_path": "tiiuae/falcon-40b",
+  "alibi": false,
+  "apply_residual_connection_post_layernorm": false,
+  "architectures": [
+    "RWForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "tiiuae/falcon-40b--configuration_RW.RWConfig",
+    "AutoModel": "tiiuae/falcon-40b--modelling_RW.RWModel",
+    "AutoModelForCausalLM": "tiiuae/falcon-40b--modelling_RW.RWForCausalLM",
+    "AutoModelForQuestionAnswering": "tiiuae/falcon-40b--modelling_RW.RWForQuestionAnswering",
+    "AutoModelForSequenceClassification": "tiiuae/falcon-40b--modelling_RW.RWForSequenceClassification",
+    "AutoModelForTokenClassification": "tiiuae/falcon-40b--modelling_RW.RWForTokenClassification"
+  },
+  "bias": false,
+  "bos_token_id": 11,
+  "custom_pipelines": {
+    "text-generation": {
+      "impl": "h2oai_pipeline.H2OTextGenerationPipeline",
+      "pt": "AutoModelForCausalLM"
+    }
+  },
+  "eos_token_id": 11,
+  "hidden_dropout": 0.0,
+  "hidden_size": 8192,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "RefinedWeb",
+  "n_head": 128,
+  "n_head_kv": 8,
+  "n_layer": 60,
+  "parallel_attn": true,
+  "torch_dtype": "float16",
+  "transformers_version": "4.30.0.dev0",
+  "use_cache": true,
+  "vocab_size": 65024
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "transformers_version": "4.30.0.dev0"
+}

h2oai_pipeline.py

@@ -0,0 +1,128 @@
+from transformers import TextGenerationPipeline
+from transformers.pipelines.text_generation import ReturnType
+
+from stopping import get_stopping
+from prompter import Prompter
+
+
+class H2OTextGenerationPipeline(TextGenerationPipeline):
+    def __init__(self, *args, debug=False, chat=False, stream_output=False,
+                 sanitize_bot_response=True,
+                 use_prompter=True, prompter=None, prompt_type=None,
+                 max_input_tokens=2048 - 256, **kwargs):
+        """
+        HF-like pipeline, but handle instruction prompting and stopping (for some models)
+        :param args:
+        :param debug:
+        :param chat:
+        :param stream_output:
+        :param sanitize_bot_response:
+        :param use_prompter: Whether to use prompter.  If pass prompt_type, will make prompter
+        :param prompter: prompter, can pass if have already
+        :param prompt_type: prompt_type, e.g. human_bot.  See prompt_type to model mapping in from prompter.py.
+                            If use_prompter, then will make prompter and use it.
+        :param max_input_tokens:
+        :param kwargs:
+        """
+        super().__init__(*args, **kwargs)
+        self.prompt_text = None
+        self.use_prompter = use_prompter
+        self.prompt_type = prompt_type
+        self.prompter = prompter
+        if self.use_prompter:
+            if self.prompter is not None:
+                assert self.prompter.prompt_type is not None
+            else:
+                self.prompter = Prompter(self.prompt_type, debug=debug, chat=chat, stream_output=stream_output)
+            self.human = self.prompter.humanstr
+            self.bot = self.prompter.botstr
+            self.can_stop = True
+        else:
+            self.prompter = None
+            self.human = None
+            self.bot = None
+            self.can_stop = False
+        self.sanitize_bot_response = sanitize_bot_response
+        self.max_input_tokens = max_input_tokens  # not for generate, so ok that not kwargs
+
+    def preprocess(self, prompt_text, prefix="", handle_long_generation=None, **generate_kwargs):
+        data_point = dict(context='', instruction=prompt_text, input='')
+        if self.prompter is not None:
+            prompt_text = self.prompter.generate_prompt(data_point)
+        self.prompt_text = prompt_text
+        if handle_long_generation is None:
+            # forces truncation of inputs to avoid critical failure
+            handle_long_generation = 'hole'
+        return super().preprocess(prompt_text, prefix=prefix, handle_long_generation=handle_long_generation,
+                                  **generate_kwargs)
+
+    def postprocess(self, model_outputs, return_type=ReturnType.FULL_TEXT, clean_up_tokenization_spaces=True):
+        records = super().postprocess(model_outputs, return_type=return_type,
+                                      clean_up_tokenization_spaces=clean_up_tokenization_spaces)
+        for rec in records:
+            if self.use_prompter:
+                outputs = rec['generated_text']
+                outputs = self.prompter.get_response(outputs, prompt=self.prompt_text,
+                                                     sanitize_bot_response=self.sanitize_bot_response)
+            elif self.bot and self.human:
+                outputs = rec['generated_text'].split(self.bot)[1].strip().split(self.human)[0].strip()
+            else:
+                outputs = rec['generated_text']
+            rec['generated_text'] = outputs
+        return records
+
+    def _forward(self, model_inputs, **generate_kwargs):
+        if self.can_stop:
+            stopping_criteria = get_stopping(self.prompt_type, self.tokenizer, self.device, human=self.human,
+                                             bot=self.bot)
+            generate_kwargs['stopping_criteria'] = stopping_criteria
+        # return super()._forward(model_inputs, **generate_kwargs)
+        return self.__forward(model_inputs, **generate_kwargs)
+
+    # FIXME: Copy-paste of original _forward, but removed copy.deepcopy()
+    # FIXME: https://github.com/h2oai/h2ogpt/issues/172
+    def __forward(self, model_inputs, **generate_kwargs):
+        input_ids = model_inputs["input_ids"]
+        attention_mask = model_inputs.get("attention_mask", None)
+        # Allow empty prompts
+        if input_ids.shape[1] == 0:
+            input_ids = None
+            attention_mask = None
+            in_b = 1
+        else:
+            in_b = input_ids.shape[0]
+        prompt_text = model_inputs.pop("prompt_text")
+
+        ## If there is a prefix, we may need to adjust the generation length. Do so without permanently modifying
+        ## generate_kwargs, as some of the parameterization may come from the initialization of the pipeline.
+        # generate_kwargs = copy.deepcopy(generate_kwargs)
+        prefix_length = generate_kwargs.pop("prefix_length", 0)
+        if prefix_length > 0:
+            has_max_new_tokens = "max_new_tokens" in generate_kwargs or (
+                    "generation_config" in generate_kwargs
+                    and generate_kwargs["generation_config"].max_new_tokens is not None
+            )
+            if not has_max_new_tokens:
+                generate_kwargs["max_length"] = generate_kwargs.get("max_length") or self.model.config.max_length
+                generate_kwargs["max_length"] += prefix_length
+            has_min_new_tokens = "min_new_tokens" in generate_kwargs or (
+                    "generation_config" in generate_kwargs
+                    and generate_kwargs["generation_config"].min_new_tokens is not None
+            )
+            if not has_min_new_tokens and "min_length" in generate_kwargs:
+                generate_kwargs["min_length"] += prefix_length
+
+        # BS x SL
+        generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
+        out_b = generated_sequence.shape[0]
+        if self.framework == "pt":
+            generated_sequence = generated_sequence.reshape(in_b, out_b // in_b, *generated_sequence.shape[1:])
+        elif self.framework == "tf":
+            from transformers import is_tf_available
+            if is_tf_available():
+                import tensorflow as tf
+                generated_sequence = tf.reshape(generated_sequence,
+                                                (in_b, out_b // in_b, *generated_sequence.shape[1:]))
+            else:
+                raise ValueError("TF not avaialble.")
+        return {"generated_sequence": generated_sequence, "input_ids": input_ids, "prompt_text": prompt_text}

pytorch_model-00001-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9034a2f18ca79dc8ea947476a24e823164e055e891923732de94f3fd1d5798fa
+size 4605550559

pytorch_model-00002-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4698fe9527f7b6c2c2ca110cccf853c0f54b2e8968cba8cfb301529bd27e9bf5
+size 4899219433

pytorch_model-00003-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf29337ab7126ab311fea87df7b5a5db1512a101ef1d3ad8713cd7437d4d8bc5
+size 4613939365

pytorch_model-00004-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bad607e087ad7d6cba05fdd3218b5f664843e67dafc1ff979ad3f1674fc21b1f
+size 4899219497

pytorch_model-00005-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08d407d8f479ba2fe98fbec90d8b8803b477254b7c945e408589b084103e4436
+size 4613939429

pytorch_model-00006-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00c2dba754cf9b22cf1fdb9e4286fa2bfd10d1016efbd214861e3ed6545f3761
+size 4899219497

pytorch_model-00007-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:559beaaf574a23286509d8cd1f3f3b9ca26a6cbc0b777f2babb03a0cdf48ff9f
+size 4613939429

pytorch_model-00008-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e9a7d668ead3885e8e0dea70a54edbf486cd2f55b599101f462456fc465a6d2
+size 4899219497

pytorch_model-00009-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37e8870bae29e255c73341fae121087413b6aa3c4ad9c1af5444ce54a24b3b2b
+size 4613939429

pytorch_model-00010-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e034c07927db399d040b7801f959ed553a093058fee314d2613f1c4a0f1e01d
+size 4899219497

pytorch_model-00011-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efa7fc930145040845a5b427870af4c3e83b38d6b44cc5107beb5e0e5c583182
+size 4613939429

pytorch_model-00012-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05ba351f5935a922efc578f381b53f531a8cd04d533b0573b4be8f403c214c4d
+size 4899219497

pytorch_model-00013-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0e9a20aeeb805080b37fcdcb558367b0f002a49ad26bbcb76f7bfcb25f34dd7
+size 4613939429

pytorch_model-00014-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2d4ce817d399e23a96f9e7af7f234e828e0c4e2d7729497a75db94f9d4b0e268
+size 4899219497

pytorch_model-00015-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfdf5d88754bf688a4b181de08e2eb90c7d5ca935d011a40c285adcda36bd137
+size 4613939429

pytorch_model-00016-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e687803c5f291a6779804572836e4dc27202dc85d24c3c277e9187403adf55e
+size 4899219497

pytorch_model-00017-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe83979d7a84feee2b04ae3b2f2350f1316d9d75f701985a4382465e086524de
+size 4613939429

pytorch_model-00018-of-00018.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d89f4cb4f0db1e8ee7fe8cd2dbd395df1eeb8ac045b19c36caa3d0fa6ce32c8a
+size 1895927701

pytorch_model.bin.index.json

@@ -0,0 +1,491 @@
+{
+  "metadata": {
+    "total_size": 82606587904
+  },
+  "weight_map": {
+    "lm_head.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.0.ln_attn.bias": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.0.ln_attn.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.0.ln_mlp.bias": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.0.ln_mlp.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.0.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.0.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.0.self_attention.dense.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.0.self_attention.query_key_value.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.1.ln_attn.bias": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.1.ln_attn.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.1.ln_mlp.bias": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.1.ln_mlp.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.1.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.1.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.1.self_attention.dense.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.1.self_attention.query_key_value.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.10.ln_attn.bias": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.10.ln_attn.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.10.ln_mlp.bias": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.10.ln_mlp.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.10.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.10.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.10.self_attention.dense.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.10.self_attention.query_key_value.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.11.ln_attn.bias": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.11.ln_attn.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.11.ln_mlp.bias": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.11.ln_mlp.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.11.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.11.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.11.self_attention.dense.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.11.self_attention.query_key_value.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.12.ln_attn.bias": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.12.ln_attn.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.12.ln_mlp.bias": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.12.ln_mlp.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.12.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.12.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.12.self_attention.dense.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.12.self_attention.query_key_value.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.13.ln_attn.bias": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.13.ln_attn.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.13.ln_mlp.bias": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.13.ln_mlp.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.13.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.13.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.13.self_attention.dense.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.13.self_attention.query_key_value.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.14.ln_attn.bias": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.14.ln_attn.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.14.ln_mlp.bias": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.14.ln_mlp.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.14.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.14.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.14.self_attention.dense.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.14.self_attention.query_key_value.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.15.ln_attn.bias": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.15.ln_attn.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.15.ln_mlp.bias": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.15.ln_mlp.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.15.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.15.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.15.self_attention.dense.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.15.self_attention.query_key_value.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.16.ln_attn.bias": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.16.ln_attn.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.16.ln_mlp.bias": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.16.ln_mlp.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.16.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.16.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.16.self_attention.dense.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.16.self_attention.query_key_value.weight": "pytorch_model-00005-of-00018.bin",
+    "transformer.h.17.ln_attn.bias": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.17.ln_attn.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.17.ln_mlp.bias": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.17.ln_mlp.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.17.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.17.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.17.self_attention.dense.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.17.self_attention.query_key_value.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.18.ln_attn.bias": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.18.ln_attn.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.18.ln_mlp.bias": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.18.ln_mlp.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.18.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.18.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.18.self_attention.dense.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.18.self_attention.query_key_value.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.19.ln_attn.bias": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.19.ln_attn.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.19.ln_mlp.bias": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.19.ln_mlp.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.19.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.19.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.19.self_attention.dense.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.19.self_attention.query_key_value.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.2.ln_attn.bias": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.2.ln_attn.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.2.ln_mlp.bias": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.2.ln_mlp.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.2.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.2.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.2.self_attention.dense.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.2.self_attention.query_key_value.weight": "pytorch_model-00001-of-00018.bin",
+    "transformer.h.20.ln_attn.bias": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.20.ln_attn.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.20.ln_mlp.bias": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.20.ln_mlp.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.20.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.20.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.20.self_attention.dense.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.20.self_attention.query_key_value.weight": "pytorch_model-00006-of-00018.bin",
+    "transformer.h.21.ln_attn.bias": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.21.ln_attn.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.21.ln_mlp.bias": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.21.ln_mlp.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.21.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.21.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.21.self_attention.dense.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.21.self_attention.query_key_value.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.22.ln_attn.bias": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.22.ln_attn.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.22.ln_mlp.bias": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.22.ln_mlp.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.22.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.22.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.22.self_attention.dense.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.22.self_attention.query_key_value.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.23.ln_attn.bias": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.23.ln_attn.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.23.ln_mlp.bias": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.23.ln_mlp.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.23.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.23.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.23.self_attention.dense.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.23.self_attention.query_key_value.weight": "pytorch_model-00007-of-00018.bin",
+    "transformer.h.24.ln_attn.bias": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.24.ln_attn.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.24.ln_mlp.bias": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.24.ln_mlp.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.24.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.24.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.24.self_attention.dense.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.24.self_attention.query_key_value.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.25.ln_attn.bias": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.25.ln_attn.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.25.ln_mlp.bias": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.25.ln_mlp.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.25.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.25.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.25.self_attention.dense.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.25.self_attention.query_key_value.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.26.ln_attn.bias": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.26.ln_attn.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.26.ln_mlp.bias": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.26.ln_mlp.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.26.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.26.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.26.self_attention.dense.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.26.self_attention.query_key_value.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.27.ln_attn.bias": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.27.ln_attn.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.27.ln_mlp.bias": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.27.ln_mlp.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.27.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.27.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.27.self_attention.dense.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.27.self_attention.query_key_value.weight": "pytorch_model-00008-of-00018.bin",
+    "transformer.h.28.ln_attn.bias": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.28.ln_attn.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.28.ln_mlp.bias": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.28.ln_mlp.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.28.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.28.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.28.self_attention.dense.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.28.self_attention.query_key_value.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.29.ln_attn.bias": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.29.ln_attn.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.29.ln_mlp.bias": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.29.ln_mlp.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.29.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.29.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.29.self_attention.dense.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.29.self_attention.query_key_value.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.3.ln_attn.bias": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.3.ln_attn.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.3.ln_mlp.bias": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.3.ln_mlp.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.3.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.3.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.3.self_attention.dense.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.3.self_attention.query_key_value.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.30.ln_attn.bias": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.30.ln_attn.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.30.ln_mlp.bias": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.30.ln_mlp.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.30.mlp.dense_4h_to_h.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.30.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.30.self_attention.dense.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.30.self_attention.query_key_value.weight": "pytorch_model-00009-of-00018.bin",
+    "transformer.h.31.ln_attn.bias": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.31.ln_attn.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.31.ln_mlp.bias": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.31.ln_mlp.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.31.mlp.dense_4h_to_h.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.31.mlp.dense_h_to_4h.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.31.self_attention.dense.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.31.self_attention.query_key_value.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.32.ln_attn.bias": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.32.ln_attn.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.32.ln_mlp.bias": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.32.ln_mlp.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.32.mlp.dense_4h_to_h.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.32.mlp.dense_h_to_4h.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.32.self_attention.dense.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.32.self_attention.query_key_value.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.33.ln_attn.bias": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.33.ln_attn.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.33.ln_mlp.bias": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.33.ln_mlp.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.33.mlp.dense_4h_to_h.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.33.mlp.dense_h_to_4h.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.33.self_attention.dense.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.33.self_attention.query_key_value.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.34.ln_attn.bias": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.34.ln_attn.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.34.ln_mlp.bias": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.34.ln_mlp.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.34.mlp.dense_4h_to_h.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.34.mlp.dense_h_to_4h.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.34.self_attention.dense.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.34.self_attention.query_key_value.weight": "pytorch_model-00010-of-00018.bin",
+    "transformer.h.35.ln_attn.bias": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.35.ln_attn.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.35.ln_mlp.bias": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.35.ln_mlp.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.35.mlp.dense_4h_to_h.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.35.mlp.dense_h_to_4h.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.35.self_attention.dense.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.35.self_attention.query_key_value.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.36.ln_attn.bias": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.36.ln_attn.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.36.ln_mlp.bias": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.36.ln_mlp.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.36.mlp.dense_4h_to_h.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.36.mlp.dense_h_to_4h.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.36.self_attention.dense.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.36.self_attention.query_key_value.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.37.ln_attn.bias": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.37.ln_attn.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.37.ln_mlp.bias": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.37.ln_mlp.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.37.mlp.dense_4h_to_h.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.37.mlp.dense_h_to_4h.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.37.self_attention.dense.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.37.self_attention.query_key_value.weight": "pytorch_model-00011-of-00018.bin",
+    "transformer.h.38.ln_attn.bias": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.38.ln_attn.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.38.ln_mlp.bias": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.38.ln_mlp.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.38.mlp.dense_4h_to_h.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.38.mlp.dense_h_to_4h.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.38.self_attention.dense.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.38.self_attention.query_key_value.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.39.ln_attn.bias": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.39.ln_attn.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.39.ln_mlp.bias": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.39.ln_mlp.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.39.mlp.dense_4h_to_h.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.39.mlp.dense_h_to_4h.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.39.self_attention.dense.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.39.self_attention.query_key_value.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.4.ln_attn.bias": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.4.ln_attn.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.4.ln_mlp.bias": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.4.ln_mlp.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.4.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.4.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.4.self_attention.dense.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.4.self_attention.query_key_value.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.40.ln_attn.bias": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.40.ln_attn.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.40.ln_mlp.bias": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.40.ln_mlp.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.40.mlp.dense_4h_to_h.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.40.mlp.dense_h_to_4h.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.40.self_attention.dense.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.40.self_attention.query_key_value.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.41.ln_attn.bias": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.41.ln_attn.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.41.ln_mlp.bias": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.41.ln_mlp.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.41.mlp.dense_4h_to_h.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.41.mlp.dense_h_to_4h.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.41.self_attention.dense.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.41.self_attention.query_key_value.weight": "pytorch_model-00012-of-00018.bin",
+    "transformer.h.42.ln_attn.bias": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.42.ln_attn.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.42.ln_mlp.bias": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.42.ln_mlp.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.42.mlp.dense_4h_to_h.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.42.mlp.dense_h_to_4h.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.42.self_attention.dense.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.42.self_attention.query_key_value.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.43.ln_attn.bias": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.43.ln_attn.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.43.ln_mlp.bias": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.43.ln_mlp.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.43.mlp.dense_4h_to_h.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.43.mlp.dense_h_to_4h.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.43.self_attention.dense.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.43.self_attention.query_key_value.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.44.ln_attn.bias": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.44.ln_attn.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.44.ln_mlp.bias": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.44.ln_mlp.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.44.mlp.dense_4h_to_h.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.44.mlp.dense_h_to_4h.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.44.self_attention.dense.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.44.self_attention.query_key_value.weight": "pytorch_model-00013-of-00018.bin",
+    "transformer.h.45.ln_attn.bias": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.45.ln_attn.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.45.ln_mlp.bias": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.45.ln_mlp.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.45.mlp.dense_4h_to_h.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.45.mlp.dense_h_to_4h.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.45.self_attention.dense.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.45.self_attention.query_key_value.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.46.ln_attn.bias": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.46.ln_attn.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.46.ln_mlp.bias": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.46.ln_mlp.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.46.mlp.dense_4h_to_h.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.46.mlp.dense_h_to_4h.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.46.self_attention.dense.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.46.self_attention.query_key_value.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.47.ln_attn.bias": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.47.ln_attn.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.47.ln_mlp.bias": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.47.ln_mlp.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.47.mlp.dense_4h_to_h.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.47.mlp.dense_h_to_4h.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.47.self_attention.dense.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.47.self_attention.query_key_value.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.48.ln_attn.bias": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.48.ln_attn.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.48.ln_mlp.bias": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.48.ln_mlp.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.48.mlp.dense_4h_to_h.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.48.mlp.dense_h_to_4h.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.48.self_attention.dense.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.48.self_attention.query_key_value.weight": "pytorch_model-00014-of-00018.bin",
+    "transformer.h.49.ln_attn.bias": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.49.ln_attn.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.49.ln_mlp.bias": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.49.ln_mlp.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.49.mlp.dense_4h_to_h.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.49.mlp.dense_h_to_4h.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.49.self_attention.dense.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.49.self_attention.query_key_value.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.5.ln_attn.bias": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.5.ln_attn.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.5.ln_mlp.bias": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.5.ln_mlp.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.5.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.5.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.5.self_attention.dense.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.5.self_attention.query_key_value.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.50.ln_attn.bias": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.50.ln_attn.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.50.ln_mlp.bias": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.50.ln_mlp.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.50.mlp.dense_4h_to_h.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.50.mlp.dense_h_to_4h.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.50.self_attention.dense.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.50.self_attention.query_key_value.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.51.ln_attn.bias": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.51.ln_attn.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.51.ln_mlp.bias": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.51.ln_mlp.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.51.mlp.dense_4h_to_h.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.51.mlp.dense_h_to_4h.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.51.self_attention.dense.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.51.self_attention.query_key_value.weight": "pytorch_model-00015-of-00018.bin",
+    "transformer.h.52.ln_attn.bias": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.52.ln_attn.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.52.ln_mlp.bias": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.52.ln_mlp.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.52.mlp.dense_4h_to_h.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.52.mlp.dense_h_to_4h.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.52.self_attention.dense.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.52.self_attention.query_key_value.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.53.ln_attn.bias": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.53.ln_attn.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.53.ln_mlp.bias": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.53.ln_mlp.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.53.mlp.dense_4h_to_h.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.53.mlp.dense_h_to_4h.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.53.self_attention.dense.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.53.self_attention.query_key_value.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.54.ln_attn.bias": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.54.ln_attn.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.54.ln_mlp.bias": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.54.ln_mlp.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.54.mlp.dense_4h_to_h.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.54.mlp.dense_h_to_4h.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.54.self_attention.dense.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.54.self_attention.query_key_value.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.55.ln_attn.bias": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.55.ln_attn.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.55.ln_mlp.bias": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.55.ln_mlp.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.55.mlp.dense_4h_to_h.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.55.mlp.dense_h_to_4h.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.55.self_attention.dense.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.55.self_attention.query_key_value.weight": "pytorch_model-00016-of-00018.bin",
+    "transformer.h.56.ln_attn.bias": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.56.ln_attn.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.56.ln_mlp.bias": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.56.ln_mlp.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.56.mlp.dense_4h_to_h.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.56.mlp.dense_h_to_4h.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.56.self_attention.dense.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.56.self_attention.query_key_value.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.57.ln_attn.bias": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.57.ln_attn.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.57.ln_mlp.bias": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.57.ln_mlp.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.57.mlp.dense_4h_to_h.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.57.mlp.dense_h_to_4h.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.57.self_attention.dense.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.57.self_attention.query_key_value.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.58.ln_attn.bias": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.58.ln_attn.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.58.ln_mlp.bias": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.58.ln_mlp.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.58.mlp.dense_4h_to_h.weight": "pytorch_model-00018-of-00018.bin",
+    "transformer.h.58.mlp.dense_h_to_4h.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.58.self_attention.dense.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.58.self_attention.query_key_value.weight": "pytorch_model-00017-of-00018.bin",
+    "transformer.h.59.ln_attn.bias": "pytorch_model-00018-of-00018.bin",
+    "transformer.h.59.ln_attn.weight": "pytorch_model-00018-of-00018.bin",
+    "transformer.h.59.ln_mlp.bias": "pytorch_model-00018-of-00018.bin",
+    "transformer.h.59.ln_mlp.weight": "pytorch_model-00018-of-00018.bin",
+    "transformer.h.59.mlp.dense_4h_to_h.weight": "pytorch_model-00018-of-00018.bin",
+    "transformer.h.59.mlp.dense_h_to_4h.weight": "pytorch_model-00018-of-00018.bin",
+    "transformer.h.59.self_attention.dense.weight": "pytorch_model-00018-of-00018.bin",
+    "transformer.h.59.self_attention.query_key_value.weight": "pytorch_model-00018-of-00018.bin",
+    "transformer.h.6.ln_attn.bias": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.6.ln_attn.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.6.ln_mlp.bias": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.6.ln_mlp.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.6.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.6.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.6.self_attention.dense.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.6.self_attention.query_key_value.weight": "pytorch_model-00002-of-00018.bin",
+    "transformer.h.7.ln_attn.bias": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.7.ln_attn.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.7.ln_mlp.bias": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.7.ln_mlp.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.7.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.7.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.7.self_attention.dense.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.7.self_attention.query_key_value.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.8.ln_attn.bias": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.8.ln_attn.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.8.ln_mlp.bias": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.8.ln_mlp.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.8.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.8.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.8.self_attention.dense.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.8.self_attention.query_key_value.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.9.ln_attn.bias": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.9.ln_attn.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.9.ln_mlp.bias": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.9.ln_mlp.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.9.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00018.bin",
+    "transformer.h.9.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.9.self_attention.dense.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.h.9.self_attention.query_key_value.weight": "pytorch_model-00003-of-00018.bin",
+    "transformer.ln_f.bias": "pytorch_model-00018-of-00018.bin",
+    "transformer.ln_f.weight": "pytorch_model-00018-of-00018.bin",
+    "transformer.word_embeddings.weight": "pytorch_model-00001-of-00018.bin"
+  }
+}

special_tokens_map.json

@@ -0,0 +1,16 @@
+{
+  "additional_special_tokens": [
+    ">>TITLE<<",
+    ">>ABSTRACT<<",
+    ">>INTRODUCTION<<",
+    ">>SUMMARY<<",
+    ">>COMMENT<<",
+    ">>ANSWER<<",
+    ">>QUESTION<<",
+    ">>DOMAIN<<",
+    ">>PREFIX<<",
+    ">>SUFFIX<<",
+    ">>MIDDLE<<"
+  ],
+  "eos_token": "<|endoftext|>"
+}

stopping.py

@@ -0,0 +1,72 @@
+import torch
+from transformers import StoppingCriteria, StoppingCriteriaList
+
+from prompter import PromptType
+
+
+class StoppingCriteriaSub(StoppingCriteria):
+
+    def __init__(self, stops=[], encounters=[], device="cuda"):
+        super().__init__()
+        assert len(stops) % len(encounters) == 0, "Number of stops and encounters must match"
+        self.encounters = encounters
+        self.stops = [stop.to(device) for stop in stops]
+        self.num_stops = [0] * len(stops)
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        for stopi, stop in enumerate(self.stops):
+            if torch.all((stop == input_ids[0][-len(stop):])).item():
+                self.num_stops[stopi] += 1
+                if self.num_stops[stopi] >= self.encounters[stopi % len(self.encounters)]:
+                    # print("Stopped", flush=True)
+                    return True
+        # print("Tokens: %s" % input_ids[0].cpu().numpy(), flush=True)
+        # print("Stop Tokens: %s" % [x.cpu().numpy() for x in self.stops], flush=True)
+        return False
+
+
+def get_stopping(prompt_type, tokenizer, device, human='<human>:', bot="<bot>:"):
+    if prompt_type in [PromptType.human_bot.name, PromptType.instruct_vicuna.name, PromptType.instruct_with_end.name]:
+        if prompt_type == PromptType.human_bot.name:
+            # encounters = [prompt.count(human) + 1, prompt.count(bot) + 1]
+            # stopping only starts once output is beyond prompt
+            # 1 human is enough to trigger, but need 2 bots, because very first view back will be bot we added
+            stop_words = [human, bot, '\n' + human, '\n' + bot]
+            encounters = [1, 2]
+        elif prompt_type == PromptType.instruct_vicuna.name:
+            # even below is not enough, generic strings and many ways to encode
+            stop_words = [
+                '### Human:',
+                """
+### Human:""",
+                """
+### Human:
+""",
+                '### Assistant:',
+                """
+### Assistant:""",
+                """
+### Assistant:
+""",
+            ]
+            encounters = [1, 2]
+        else:
+            # some instruct prompts have this as end, doesn't hurt to stop on it since not common otherwise
+            stop_words = ['### End']
+            encounters = [1]
+        stop_words_ids = [
+            tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze() for stop_word in stop_words]
+        # handle single token case
+        stop_words_ids = [x if len(x.shape) > 0 else torch.tensor([x]) for x in stop_words_ids]
+        stop_words_ids = [x for x in stop_words_ids if x.shape[0] > 0]
+        # avoid padding in front of tokens
+        if tokenizer._pad_token:  # use hidden variable to avoid annoying properly logger bug
+            stop_words_ids = [x[1:] if x[0] == tokenizer.pad_token_id and len(x) > 1 else x for x in stop_words_ids]
+        # handle fake \n added
+        stop_words_ids = [x[1:] if y[0] == '\n' else x for x, y in zip(stop_words_ids, stop_words)]
+        # build stopper
+        stopping_criteria = StoppingCriteriaList(
+            [StoppingCriteriaSub(stops=stop_words_ids, encounters=encounters, device=device)])
+    else:
+        stopping_criteria = StoppingCriteriaList()
+    return stopping_criteria

tokenizer_config.json

@@ -0,0 +1,7 @@
+{
+  "add_prefix_space": false,
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 2048,
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}