Duplicate from Muse-research/Muse2-230M

Co-authored-by: Ill Ness <Ill-Ness@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,98 @@
+---
+language:
+- en
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: text-generation
+pretty_name: Muse2-230M
+tags:
+- text-generation
+- conversational
+- hybrid-attention
+- convolution
+- grouped-query-attention
+- edge
+- on-device
+- structured-output
+- json
+- muse2
+task_categories:
+- text-generation
+- conversational
+size_categories:
+- 100M<n<1B
+base_model:
+- Muse-research/Muse2-230M-Base
+---
+
+# Muse2-230M
+
+## Model Information
+
+The Muse2 collection is a family of compact language models built **entirely from scratch** — architecture, tokenizer, and inference stack — for edge and on-device use. Muse2 pairs an LFM2-style hybrid layout (short causal convolutions + grouped-query attention) with a parallel merged SwiGLU MLP per block, chosen for inference efficiency on constrained hardware. The instruction-tuned version is optimized for assistant-like chat, formatting, and structured extraction from developer logs.
+
+**Model Developer:** Muse Research
+
+**Model Architecture:** Muse2 is an auto-regressive language model that uses a hybrid convolution/attention architecture implemented from scratch in PyTorch (`Muse2ForCausalLM`, no `transformers` code). The tuned version is aligned with supervised fine-tuning (SFT) with completion-only loss masking. No RLHF or DPO has been applied.
+
+|  | Params | Input modalities | Output modalities | Context Length | GQA | Shared Embeddings | Knowledge cutoff |
+| :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- |
+| Muse2-230M (Base)  | 230M (196.1M) | English Text | English Text and code  | 8k (128k max) | Yes | Yes | Early 2024 |
+| Muse2-230M (Instruct)  | 230M (196.1M) | English Text | English Text and code  | 8k (128k max) | Yes | Yes | Early 2024 |
+
+**Supported Languages:** English. The model also handles source code and mathematical text.
+
+**Muse2 Model Family:** All model versions use Grouped-Query Attention (GQA, 16 query / 8 key-value heads) and tied input/output embeddings. The 14-block layout interleaves 8 causal depthwise short-convolution blocks (kernel size 3) with 6 full-attention blocks; RoPE with theta 1e6; RMSNorm; parallel merged SwiGLU MLPs (hidden 1024, ff 2560); vocabulary of 65,536 byte-level BPE tokens trained from scratch (pad=0, bos=1, eos=7, ChatML-style control tokens).
+
+**Model Release Date:** Aug 23, 2026
+
+**Status:** This is a static model. Later versions may be released that improve model capabilities.
+
+**License:** Model weights are released under Apache-2.0.
+
+## Intended Use
+
+**Intended Use Cases:** Muse2 is intended for research and education: studying compact hybrid architectures, edge/on-device inference, structured log-to-JSON extraction, formatting-disciplined output, and as a base for further fine-tuning. The Instruct version is intended for assistant-like chat in resource-constrained environments where a small footprint matters more than factual reliability.
+
+**Out of Scope:** Production deployment without developer-side evaluation and guardrails. Any safety-critical, medical, legal, or factual-reliance use. Use in any manner that violates applicable laws or the upstream dataset licenses. The model has had **no safety alignment, red-teaming, or content filtering** applied.
+
+## How to use
+
+This repository uses the from-scratch `muse` package (PyTorch + safetensors only). The config schema follows the LFM2-family layout for edge-tooling compatibility, but the architecture is not a transformers-native class.
+
+```python
+import torch
+from tokenizers import Tokenizer
+from muse import Muse2ForCausalLM
+from muse.sft_data import apply_chat_template
+
+model = Muse2ForCausalLM.from_pretrained(".", dtype=torch.float32).eval()
+tk = Tokenizer.from_file("tokenizer.json")
+
+messages = [
+    {"role": "system", "content": "You are an automated, silent compiler API."},
+    {"role": "user", "content": "<raw log> + <required schema>"},
+]
+ids = apply_chat_template(tk, messages)  # ends at the assistant header
+out = model.generate(torch.tensor([ids]), max_new_tokens=256,
+                     temperature=0.2, top_k=50, eos_token_id=4)  # <|im_end|>
+print(tk.decode(out[0][len(ids):].tolist(), skip_special_tokens=True))
+```
+
+Interactive chat on a single cloud T4:
+
+```bash
+modal run modal_app.py::chat
+```
+
+Local CPU chat: `python scripts/chat.py --model . --interactive`
+
+## Responsibility & Safety
+
+Muse2 is a research artifact released without safety fine-tuning, red-teaming, or content classification. It can produce inaccurate, biased, repetitive, or otherwise objectionable output, and it **hallucinates facts readily at this model scale**. Developers are solely responsible for deployment decisions and should add appropriate input/output safeguards for any application, and should not rely on the model for correctness.
+
+## Ethical Considerations and Limitations
+
+**Values:** Muse2 is intended to make small-scale, from-scratch language-model development accessible — every component (tokenizer, architecture, inference stack) is open and reproducible.
+
+**Testing:** Safety and capability testing to date is limited to qualitative inspection. It does not cover, nor could it cover, all scenarios. Before any deployment, developers should perform testing appropriate to their use case. Upstream dataset licenses and attribution requirements continue to apply to derivative uses.

config.json

@@ -0,0 +1,66 @@
+{
+  "architectures": [
+    "Muse2ForCausalLM"
+  ],
+  "block__name_mlp": "parallel_mlp_merged",
+  "block_auto_adjust_ff_dim": false,
+  "block_dim": 1024,
+  "block_ff_dim": 2560,
+  "block_ffn_dim_multiplier": 1.0,
+  "block_ffn_te_autocast": false,
+  "block_ffn_use_quantized_params": false,
+  "block_mlp_init_scale": 1.0,
+  "block_multiple_of": 256,
+  "block_norm_eps": 1e-05,
+  "block_out_init_scale": 1.0,
+  "block_sequence_parallel_norm_across_tp": false,
+  "block_use_swiglu": true,
+  "block_use_xavier_init": true,
+  "bos_token_id": 1,
+  "conv_L_cache": 3,
+  "conv_bias": false,
+  "conv_dim": 1024,
+  "conv_use_xavier_init": true,
+  "dtype": "bfloat16",
+  "eos_token_id": 7,
+  "ffn_te_autocast": false,
+  "ffn_use_quantized_params": false,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 2560,
+  "layer_types": [
+    "conv",
+    "conv",
+    "full_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "full_attention",
+    "conv",
+    "full_attention",
+    "conv"
+  ],
+  "max_position_embeddings": 128000,
+  "model_type": "muse2",
+  "norm_eps": 1e-05,
+  "num_attention_heads": 16,
+  "num_heads": 16,
+  "num_hidden_layers": 14,
+  "num_key_value_heads": 8,
+  "pad_token_id": 0,
+  "rope_parameters": {
+    "rope_theta": 1000000.0,
+    "rope_type": "default"
+  },
+  "sequence_parallel_norm_across_tp": false,
+  "tie_embedding": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.2.0",
+  "use_cache": true,
+  "use_pos_enc": true,
+  "vocab_size": 65536
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": 4,
+  "pad_token_id": 0,
+  "temperature": 0.8,
+  "top_k": 50,
+  "top_p": 0.95,
+  "transformers_version": "5.2.0"
+}

model.safetensors

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

special_tokens_map.json

@@ -0,0 +1,10 @@
+{
+  "<|bos|>": "<|bos|>",
+  "<|eos|>": "<|eos|>",
+  "<|im_end|>": "<|im_end|>",
+  "<|im_start|>": "<|im_start|>",
+  "<|pad|>": "<|pad|>",
+  "<|system|>": "<|system|>",
+  "<|unk|>": "<|unk|>",
+  "<|user|>": "<|user|>"
+}

tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "bos_token": "<|bos|>",
+  "bos_token_id": 1,
+  "chat_template": "{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}",
+  "eos_token": "<|eos|>",
+  "eos_token_id": 7,
+  "model_max_length": 128000,
+  "pad_token": "<|pad|>",
+  "pad_token_id": 0,
+  "unk_token": "<|unk|>",
+  "vocab_size": 65536
+}