add

.gitignore

@@ -0,0 +1 @@
+__pycache__

__init__.py

@@ -1,31 +1,31 @@
-from .configuration_nanotts import NanoTTSConfig
-from .modeling_nanotts_global_local import (
-    NanoTTSGenerationOutput,
-    NanoTTSGlobalLocalForCausalLM,
-    NanoTTSOutput,
+from .configuration_moss_tts_nano import MossTTSNanoConfig
+from .modeling_moss_tts_nano import (
+    MossTTSNanoForCausalLM,
+    MossTTSNanoGenerationOutput,
+    MossTTSNanoOutput,
 )
-from .tokenization_nanotts_sentencepiece import NanoTTSSentencePieceTokenizer
+from .tokenization_moss_tts_nano import MossTTSNanoSentencePieceTokenizer
 
 try:
-    NanoTTSConfig.register_for_auto_class()
+    MossTTSNanoConfig.register_for_auto_class()
 except Exception:
     pass
 
 for auto_class_name in ("AutoModel", "AutoModelForCausalLM"):
     try:
-        NanoTTSGlobalLocalForCausalLM.register_for_auto_class(auto_class_name)
+        MossTTSNanoForCausalLM.register_for_auto_class(auto_class_name)
     except Exception:
         pass
 
 try:
-    NanoTTSSentencePieceTokenizer.register_for_auto_class("AutoTokenizer")
+    MossTTSNanoSentencePieceTokenizer.register_for_auto_class("AutoTokenizer")
 except Exception:
     pass
 
 __all__ = [
-    "NanoTTSConfig",
-    "NanoTTSGlobalLocalForCausalLM",
-    "NanoTTSSentencePieceTokenizer",
-    "NanoTTSGenerationOutput",
-    "NanoTTSOutput",
+    "MossTTSNanoConfig",
+    "MossTTSNanoForCausalLM",
+    "MossTTSNanoSentencePieceTokenizer",
+    "MossTTSNanoGenerationOutput",
+    "MossTTSNanoOutput",
 ]

config.json

@@ -1,7 +1,7 @@
 {
   "add_cross_attention": false,
   "architectures": [
-    "NanoTTSGlobalLocalForCausalLM"
+    "MossTTSNanoForCausalLM"
   ],
   "attn_implementation": "flash_attention_2",
   "audio_assistant_slot_token_id": 9,
@@ -156,7 +156,7 @@
   "max_position_embeddings": 32768,
   "min_length": 0,
   "model_architecture": "global_local_transformer",
-  "model_type": "nano_tts",
+  "model_type": "moss_tts_nano",
   "n_vq": 16,
   "no_repeat_ngram_size": 0,
   "num_beam_groups": 1,
@@ -180,7 +180,7 @@
   "tf_legacy_loss": false,
   "tie_encoder_decoder": false,
   "tie_word_embeddings": true,
-  "tokenizer_class": "NanoTTSSentencePieceTokenizer",
+  "tokenizer_class": "MossTTSNanoSentencePieceTokenizer",
   "tokenizer_use_fast": false,
   "top_k": 50,
   "top_p": 1.0,
@@ -190,8 +190,8 @@
   "use_bfloat16": false,
   "vocab_size": 16384,
   "auto_map": {
-    "AutoConfig": "configuration_nanotts.NanoTTSConfig",
-    "AutoModel": "modeling_nanotts_global_local.NanoTTSGlobalLocalForCausalLM",
-    "AutoModelForCausalLM": "modeling_nanotts_global_local.NanoTTSGlobalLocalForCausalLM"
+    "AutoConfig": "configuration_moss_tts_nano.MossTTSNanoConfig",
+    "AutoModel": "modeling_moss_tts_nano.MossTTSNanoForCausalLM",
+    "AutoModelForCausalLM": "modeling_moss_tts_nano.MossTTSNanoForCausalLM"
   }
 }

configuration_nanotts.py → configuration_moss_tts_nano.py

@@ -5,8 +5,8 @@ from transformers.configuration_utils import PretrainedConfig
 from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 
 
-class NanoTTSConfig(PretrainedConfig):
-    model_type = "nano_tts"
+class MossTTSNanoConfig(PretrainedConfig):
+    model_type = "moss_tts_nano"
     keys_to_ignore_at_inference = ["past_key_values"]
 
     def __init__(
@@ -103,3 +103,6 @@ class NanoTTSConfig(PretrainedConfig):
         output = super().to_dict()
         output["gpt2_config"] = self.gpt2_config.to_dict()
         return output
+
+
+__all__ = ["MossTTSNanoConfig"]

gpt2_decoder.py

@@ -33,7 +33,7 @@ class PackedSequenceMetadata:
     seq_len: Optional[int] = None
 
 
-class NanoGPT2RotaryEmbedding(nn.Module):
+class MossTTSNanoGPT2RotaryEmbedding(nn.Module):
     def __init__(self, dim: int, base: float = 10000.0) -> None:
         super().__init__()
         if dim % 2 != 0:
@@ -70,7 +70,7 @@ def apply_rotary_pos_emb(
     return (hidden_states * cos) + (rotate_half(hidden_states) * sin)
 
 
-class NanoGPT2MLP(nn.Module):
+class MossTTSNanoGPT2MLP(nn.Module):
     def __init__(self, config: GPT2Config) -> None:
         super().__init__()
         hidden_size = int(config.hidden_size)
@@ -87,7 +87,7 @@ class NanoGPT2MLP(nn.Module):
         return self.dropout(hidden_states)
 
 
-class NanoGPT2Attention(nn.Module):
+class MossTTSNanoGPT2Attention(nn.Module):
     def __init__(self, config: GPT2Config, layer_idx: int, attn_implementation: str) -> None:
         super().__init__()
         hidden_size = int(config.hidden_size)
@@ -112,7 +112,7 @@ class NanoGPT2Attention(nn.Module):
         self.c_proj = nn.Linear(hidden_size, hidden_size)
         self.rotary_emb = None
         if self.position_embedding_type == "rope":
-            self.rotary_emb = NanoGPT2RotaryEmbedding(
+            self.rotary_emb = MossTTSNanoGPT2RotaryEmbedding(
                 self.head_dim,
                 base=float(getattr(config, "rope_base", 10000.0)),
             )
@@ -353,14 +353,14 @@ class NanoGPT2Attention(nn.Module):
         return self.resid_dropout(attn_output), present
 
 
-class NanoGPT2Block(nn.Module):
+class MossTTSNanoGPT2Block(nn.Module):
     def __init__(self, config: GPT2Config, layer_idx: int, attn_implementation: str) -> None:
         super().__init__()
         hidden_size = int(config.hidden_size)
         self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-        self.attn = NanoGPT2Attention(config, layer_idx=layer_idx, attn_implementation=attn_implementation)
+        self.attn = MossTTSNanoGPT2Attention(config, layer_idx=layer_idx, attn_implementation=attn_implementation)
         self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
-        self.mlp = NanoGPT2MLP(config)
+        self.mlp = MossTTSNanoGPT2MLP(config)
 
     def forward(
         self,
@@ -384,7 +384,7 @@ class NanoGPT2Block(nn.Module):
         return hidden_states, present
 
 
-class NanoGPT2Model(nn.Module):
+class MossTTSNanoGPT2Model(nn.Module):
     def __init__(self, config: GPT2Config, attn_implementation: str = "eager") -> None:
         super().__init__()
         self.config = config
@@ -397,7 +397,7 @@ class NanoGPT2Model(nn.Module):
         self.wpe = nn.Embedding(config.n_positions, hidden_size) if self.position_embedding_type == "absolute" else nn.Identity()
         self.drop = nn.Dropout(config.embd_pdrop)
         self.h = nn.ModuleList(
-            [NanoGPT2Block(config, layer_idx=index, attn_implementation=attn_implementation) for index in range(config.n_layer)]
+            [MossTTSNanoGPT2Block(config, layer_idx=index, attn_implementation=attn_implementation) for index in range(config.n_layer)]
         )
         self.ln_f = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
         self.gradient_checkpointing = False
@@ -445,7 +445,7 @@ class NanoGPT2Model(nn.Module):
         batch_size, seq_len = cu_seqlens.shape[0], cu_seqlens.shape[1] - 1
         device = cu_seqlens.device
         position_ids = torch.zeros((batch_size, seq_len), dtype=torch.long, device=device)
-        counts = NanoGPT2Model._normalize_num_sequences(cu_seqlens, num_sequences, device=device)
+        counts = MossTTSNanoGPT2Model._normalize_num_sequences(cu_seqlens, num_sequences, device=device)
         for batch_index in range(batch_size):
             sequence_count = int(counts[batch_index].item())
             boundaries = cu_seqlens[batch_index, : sequence_count + 1].tolist()
@@ -467,7 +467,7 @@ class NanoGPT2Model(nn.Module):
         if cu_seqlens.ndim == 1:
             cu_seqlens = cu_seqlens.unsqueeze(0)
         device = hidden_states.device
-        counts = NanoGPT2Model._normalize_num_sequences(cu_seqlens, num_sequences, device=device)
+        counts = MossTTSNanoGPT2Model._normalize_num_sequences(cu_seqlens, num_sequences, device=device)
         flat_indices = []
         cumulative = [0]
         max_seqlen = 0

modeling_nanotts_global_local.py → modeling_moss_tts_nano.py

@@ -19,19 +19,19 @@ from transformers.modeling_outputs import ModelOutput
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 
-from .configuration_nanotts import NanoTTSConfig
-from .gpt2_decoder import NanoGPT2Block, NanoGPT2Model
+from .configuration_moss_tts_nano import MossTTSNanoConfig
+from .gpt2_decoder import MossTTSNanoGPT2Block, MossTTSNanoGPT2Model
 from .prompting import (
     build_assistant_prompt_prefix,
     build_prompt_token_ids,
     build_user_prompt_after_reference,
     build_user_prompt_prefix,
 )
-from .tokenization_nanotts_sentencepiece import NanoTTSSentencePieceTokenizer
+from .tokenization_moss_tts_nano import MossTTSNanoSentencePieceTokenizer
 
 
 @dataclass
-class NanoTTSOutput(ModelOutput):
+class MossTTSNanoOutput(ModelOutput):
     global_hidden_states: Optional[torch.FloatTensor] = None
     past_key_values: Optional[tuple[tuple[torch.Tensor, torch.Tensor], ...]] = None
     hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
@@ -39,7 +39,7 @@ class NanoTTSOutput(ModelOutput):
 
 
 @dataclass
-class NanoTTSGenerationOutput(ModelOutput):
+class MossTTSNanoGenerationOutput(ModelOutput):
     audio_token_ids: torch.LongTensor
     prompt_input_ids: Optional[torch.LongTensor] = None
 
@@ -55,24 +55,24 @@ _CLAUSE_SPLIT_PUNCTUATION = frozenset(",，、；;：:")
 _CLOSING_PUNCTUATION = frozenset("\"'”’)]}）】》」』")
 
 
-class NanoTTSPreTrainedModel(PreTrainedModel):
-    config_class = NanoTTSConfig
+class MossTTSNanoPreTrainedModel(PreTrainedModel):
+    config_class = MossTTSNanoConfig
     base_model_prefix = "transformer"
     supports_gradient_checkpointing = False
-    _no_split_modules = ["NanoGPT2Block"]
+    _no_split_modules = ["MossTTSNanoGPT2Block"]
     _supports_flash_attn_2 = True
     _supports_sdpa = True
 
 
-class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
+class MossTTSNanoForCausalLM(MossTTSNanoPreTrainedModel):
     _keys_to_ignore_on_load_unexpected = [r"local_transformer\.wte\.weight"]
 
-    def __init__(self, config: NanoTTSConfig) -> None:
+    def __init__(self, config: MossTTSNanoConfig) -> None:
         super().__init__(config)
         config.gpt2_config.pad_token_id = config.pad_token_id
         config.gpt2_config._attn_implementation = config.attn_implementation
 
-        self.transformer = NanoGPT2Model(
+        self.transformer = MossTTSNanoGPT2Model(
             config.gpt2_config,
             attn_implementation=config.attn_implementation,
         )
@@ -97,7 +97,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         local_gpt2_config["n_layer"] = int(config.local_transformer_layers)
         local_gpt2_config["n_positions"] = config.n_vq + 1
         local_gpt2_config["n_ctx"] = config.n_vq + 1
-        self.local_transformer = NanoGPT2Model(
+        self.local_transformer = MossTTSNanoGPT2Model(
             GPT2Config(**local_gpt2_config),
             attn_implementation=str(config.local_transformer_attn_implementation),
         )
@@ -210,7 +210,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
                 outputs.attentions,
             )
 
-        return NanoTTSOutput(
+        return MossTTSNanoOutput(
             global_hidden_states=outputs.last_hidden_state,
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
@@ -258,7 +258,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         while "  " in normalized_text:
             normalized_text = normalized_text.replace("  ", " ")
 
-        contains_cjk = NanoTTSGlobalLocalForCausalLM._contains_cjk(normalized_text)
+        contains_cjk = MossTTSNanoForCausalLM._contains_cjk(normalized_text)
         if contains_cjk:
             if normalized_text[-1] not in _SENTENCE_END_PUNCTUATION:
                 normalized_text = normalized_text + "。"
@@ -360,7 +360,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
             return right
         if not right:
             return left
-        if NanoTTSGlobalLocalForCausalLM._contains_cjk(left) or NanoTTSGlobalLocalForCausalLM._contains_cjk(right):
+        if MossTTSNanoForCausalLM._contains_cjk(left) or MossTTSNanoForCausalLM._contains_cjk(right):
             return left + right
         return left + " " + right
 
@@ -880,7 +880,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         sample_rate: int,
         first_audio_emitted_at: Optional[float],
     ) -> int:
-        lead_seconds = NanoTTSGlobalLocalForCausalLM._compute_stream_lead_seconds(
+        lead_seconds = MossTTSNanoForCausalLM._compute_stream_lead_seconds(
             emitted_samples_total=emitted_samples_total,
             sample_rate=sample_rate,
             first_audio_emitted_at=first_audio_emitted_at,
@@ -1040,7 +1040,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
             hf_dir = candidate_path / "hf_tokenizer"
             if self._looks_like_hf_tokenizer_dir(hf_dir):
                 return hf_dir
-            sentencepiece_model = candidate_path / "sentencepiece" / "nanotts_spm_bpe.model"
+            sentencepiece_model = candidate_path / "sentencepiece" / "mossttsnano_spm_bpe.model"
             if sentencepiece_model.is_file():
                 return sentencepiece_model
             final_summary_path = candidate_path / "final_summary.json"
@@ -1057,7 +1057,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
 
     def _load_resolved_text_tokenizer(self, resolved_path: Path, cache_dir: str):
         if resolved_path.is_file() and resolved_path.suffix == ".model":
-            return NanoTTSSentencePieceTokenizer(vocab_file=str(resolved_path))
+            return MossTTSNanoSentencePieceTokenizer(vocab_file=str(resolved_path))
         try:
             return AutoTokenizer.from_pretrained(
                 str(resolved_path),
@@ -1069,7 +1069,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         except Exception:
             model_path = resolved_path / "tokenizer.model"
             if model_path.is_file():
-                return NanoTTSSentencePieceTokenizer(vocab_file=str(model_path))
+                return MossTTSNanoSentencePieceTokenizer(vocab_file=str(model_path))
             raise
 
     @staticmethod
@@ -1276,7 +1276,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
         if isinstance(output, (list, tuple)) and output:
             if len(output) == 2 and isinstance(output[1], (int, float)):
                 return output[0]
-            return NanoTTSGlobalLocalForCausalLM._extract_tensor_candidate(output[0])
+            return MossTTSNanoForCausalLM._extract_tensor_candidate(output[0])
         raise TypeError(f"Unsupported audio tokenizer output type: {type(output)!r}")
 
     @staticmethod
@@ -1750,7 +1750,7 @@ class NanoTTSGlobalLocalForCausalLM(NanoTTSPreTrainedModel):
             return
         yield {
             "type": "final",
-            "generation": NanoTTSGenerationOutput(audio_token_ids=audio_token_ids, prompt_input_ids=input_ids),
+            "generation": MossTTSNanoGenerationOutput(audio_token_ids=audio_token_ids, prompt_input_ids=input_ids),
         }
 
     @torch.no_grad()

prompting.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import List, Sequence
 
-from .configuration_nanotts import NanoTTSConfig
+from .configuration_moss_tts_nano import MossTTSNanoConfig
 
 
 USER_ROLE_PREFIX = "user\n"
@@ -48,7 +48,7 @@ def decode_text(tokenizer, token_ids: Sequence[int]) -> str:
             return str(tokenizer.decode(list(token_ids)))
 
 
-def build_user_prompt_prefix(tokenizer, config: NanoTTSConfig) -> List[int]:
+def build_user_prompt_prefix(tokenizer, config: MossTTSNanoConfig) -> List[int]:
     return [config.im_start_token_id] + encode_text(tokenizer, USER_ROLE_PREFIX) + encode_text(
         tokenizer,
         USER_TEMPLATE_REFERENCE_PREFIX,
@@ -59,7 +59,7 @@ def build_user_prompt_after_reference(tokenizer) -> List[int]:
     return encode_text(tokenizer, USER_TEMPLATE_AFTER_REFERENCE)
 
 
-def build_assistant_prompt_prefix(tokenizer, config: NanoTTSConfig) -> List[int]:
+def build_assistant_prompt_prefix(tokenizer, config: MossTTSNanoConfig) -> List[int]:
     return encode_text(tokenizer, USER_TEMPLATE_SUFFIX) + [config.im_end_token_id] + encode_text(
         tokenizer,
         ASSISTANT_TURN_PREFIX,
@@ -69,7 +69,7 @@ def build_assistant_prompt_prefix(tokenizer, config: NanoTTSConfig) -> List[int]
     )
 
 
-def build_prompt_prefix(tokenizer, config: NanoTTSConfig) -> List[int]:
+def build_prompt_prefix(tokenizer, config: MossTTSNanoConfig) -> List[int]:
     return (
         build_user_prompt_prefix(tokenizer, config)
         + encode_text(tokenizer, "None")
@@ -77,13 +77,13 @@ def build_prompt_prefix(tokenizer, config: NanoTTSConfig) -> List[int]:
     )
 
 
-def build_prompt_suffix(tokenizer, config: NanoTTSConfig) -> List[int]:
+def build_prompt_suffix(tokenizer, config: MossTTSNanoConfig) -> List[int]:
     return build_assistant_prompt_prefix(tokenizer, config)
 
 
 def build_prompt_token_ids(
     tokenizer,
-    config: NanoTTSConfig,
+    config: MossTTSNanoConfig,
     text_token_ids: Sequence[int],
 ) -> List[int]:
     return build_prompt_prefix(tokenizer, config) + [int(token_id) for token_id in text_token_ids] + build_prompt_suffix(

tokenization_nanotts_sentencepiece.py → tokenization_moss_tts_nano.py

@@ -11,7 +11,7 @@ from transformers import PreTrainedTokenizer
 VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
 
 
-class NanoTTSSentencePieceTokenizer(PreTrainedTokenizer):
+class MossTTSNanoSentencePieceTokenizer(PreTrainedTokenizer):
     vocab_files_names = VOCAB_FILES_NAMES
     model_input_names = ["input_ids", "attention_mask"]
 
@@ -101,3 +101,6 @@ class NanoTTSSentencePieceTokenizer(PreTrainedTokenizer):
         if token_ids_1 is None:
             return [0] * len(token_ids_0)
         return [0] * (len(token_ids_0) + len(token_ids_1))
+
+
+__all__ = ["MossTTSNanoSentencePieceTokenizer"]

tokenizer_config.json

@@ -36,7 +36,7 @@
   "additional_special_tokens": [],
   "auto_map": {
     "AutoTokenizer": [
-      "tokenization_nanotts_sentencepiece.NanoTTSSentencePieceTokenizer",
+      "tokenization_moss_tts_nano.MossTTSNanoSentencePieceTokenizer",
       null
     ]
   },
@@ -47,6 +47,6 @@
   "extra_special_tokens": {},
   "model_max_length": 16384,
   "pad_token": "<pad>",
-  "tokenizer_class": "NanoTTSSentencePieceTokenizer",
+  "tokenizer_class": "MossTTSNanoSentencePieceTokenizer",
   "unk_token": "<unk>"
 }