fix: Restructure and rewite sampling logic to be compatible with split. (#587)

* Restructure and rewite sampling logic to be compatible with split.

Signed-off-by: Dushyant Behl <[email protected]>

* Restructure and rewite sampling logic to be compatible with split.

Signed-off-by: Dushyant Behl <[email protected]>

* Add a cached maykeye/tinyllama due to HF rate limiting crashing tests.

Signed-off-by: Dushyant Behl <[email protected]>

---------

Signed-off-by: Dushyant Behl <[email protected]>

Author: dushyantbehl

tests/artifacts/language_models/__init__.py

@@ -0,0 +1,22 @@
+# Copyright The FMS HF Tuning Authors
+#
+# 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.
+
+"""Helpful cached languaged models for unit tests.
+"""
+# Standard
+import os
+
+### Constants used for model path
+PREDEFINED_MODEL_PATH = os.path.join(os.path.dirname(__file__))
+MAYKEYE_TINY_LLAMA_CACHED = os.path.join(PREDEFINED_MODEL_PATH, "maykeye-tinyllama-v0")

tests/artifacts/language_models/maykeye-tinyllama-v0/config.json

@@ -0,0 +1,22 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_act": "silu",
+  "hidden_size": 64,
+  "initializer_range": 0.02,
+  "intermediate_size": 256,
+  "max_position_embeddings": 2048,
+  "model_type": "llama",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 8,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.30.2",
+  "use_cache": true,
+  "vocab_size": 32000
+}

tests/artifacts/language_models/maykeye-tinyllama-v0/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "transformers_version": "4.30.2"
+}

tests/artifacts/language_models/maykeye-tinyllama-v0/model.safetensors

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}