Add S2-Pro finetune part.

fish_speech/configs/lora/r_8_alpha_16.yaml

@@ -2,3 +2,6 @@ _target_: fish_speech.models.text2semantic.lora.LoraConfig
 r: 8
 lora_alpha: 16
 lora_dropout: 0.01
+target_modules:
+  - attention
+  - mlp

fish_speech/configs/text2semantic_finetune.yaml

@@ -12,12 +12,18 @@ trainer:
   gradient_clip_val: 1.0
   gradient_clip_algorithm: "norm"
   max_steps: 10000
-  precision: bf16-true
+  precision: bf16-mixed
   limit_val_batches: 10
   val_check_interval: 100
   # strategy:
   #   find_unused_parameters: true
-  #   static_graph: true 
+  #   static_graph: true
+
+loss:
+  base_weight: 1.0
+  base_vq_weight: 0.5
+  decode_semantic_token_weight: 0.5
+  semantic_weights: null
 
 # Dataset Configuration
 tokenizer:
@@ -28,7 +34,7 @@ tokenizer:
 train_dataset:
   _target_: fish_speech.datasets.semantic.AutoTextSemanticInstructionIterableDataset
   proto_files:
-    - data/protos
+    - data/quantized-dataset-ft
   tokenizer: ${tokenizer}
   causal: true
   max_length: ${max_length}
@@ -38,7 +44,7 @@ train_dataset:
 val_dataset:
   _target_: fish_speech.datasets.semantic.AutoTextSemanticInstructionIterableDataset
   proto_files:
-    - data/protos
+    - data/quantized-dataset-ft
   tokenizer: ${tokenizer}
   causal: true
   max_length: ${max_length}
@@ -57,28 +63,33 @@ data:
 # Model Configuration
 model:
   _target_: fish_speech.models.text2semantic.lit_module.TextToSemantic
-  model: 
+  model:
     _target_: fish_speech.models.text2semantic.llama.BaseTransformer.from_pretrained
     path: ${pretrained_ckpt_path}
     load_weights: true
     max_length: ${max_length}
     lora_config: null
 
+  base_weight: ${loss.base_weight}
+  base_vq_weight: ${loss.base_vq_weight}
+  decode_semantic_token_weight: ${loss.decode_semantic_token_weight}
+  semantic_weights: ${loss.semantic_weights}
+
   optimizer:
     _target_: torch.optim.AdamW
     _partial_: true
-    lr: 1e-4
+    lr: 2e-5
     weight_decay: 0
     betas: [0.9, 0.95]
-    eps: 1e-5
+    eps: 1e-6
 
   lr_scheduler:
     _target_: torch.optim.lr_scheduler.LambdaLR
     _partial_: true
     lr_lambda:
       _target_: fish_speech.scheduler.get_constant_schedule_with_warmup_lr_lambda
       _partial_: true
-      num_warmup_steps: 10
+      num_warmup_steps: 100
 
 # Callbacks
 callbacks:

fish_speech/content_sequence.py

@@ -156,6 +156,7 @@ def encode(
         tokenizer: FishTokenizer,
         add_shift: bool = True,
         ignore_loss_tokens: list[str] = [],
+        max_length: int | None = None,
     ) -> EncodedMessage:
         """
         Encode the sequence parts into tokens for the model.

fish_speech/datasets/semantic.py

@@ -16,6 +16,7 @@
 from torch.utils.data import DataLoader, Dataset, IterableDataset, get_worker_info
 
 from fish_speech.content_sequence import ContentSequence, TextPart, VQPart
+from fish_speech.conversation import Conversation, Message
 
 CODEBOOK_PAD_TOKEN_ID = 0
 
@@ -190,34 +191,36 @@ def pack_sentences(
         skip_text: bool = False,
     ):
 
-        seq = ContentSequence()
-
-        seq.append(TextPart(text="Speak out the provided text."))
-
-        # User's turn
         cated_sentences = " ".join(sentences)
         if skip_text:
             cated_sentences = "<|skip_text|>"
 
-        seq.append(
-            TextPart(text=f"<|speaker:user|> {cated_sentences}"),
-            add_end=True,
-        )
+        messages = [
+            Message(
+                role="system",
+                parts=[TextPart(text="Speak out the provided text.")],
+            ),
+            Message(
+                role="user",
+                parts=[TextPart(text=cated_sentences)],
+            ),
+        ]
 
         # Assistant's turn
         vq_codes = [x.values for x in semantics[0]]
         vq_codes_tensor = torch.tensor(vq_codes).to(torch.int32)
+        vq_part = VQPart(codes=vq_codes_tensor)
 
-        # 将 cal_loss=True 直接关联到 VQPart 上，这比之前更精确
-        vq_part = VQPart(codes=vq_codes_tensor, cal_loss=True)
-
-        # 将多个 parts 一起添加，最后也加上 <|im_end|>
-        seq.append(
-            [TextPart(text="<|speaker:assistant|> <|voice|>"), vq_part],
-            add_end=True,
+        messages.append(
+            Message(
+                role="assistant",
+                parts=[TextPart(text="<|voice|>"), vq_part],
+                cal_loss=True,
+            )
         )
 
-        encoded = seq.encode(
+        conversation = Conversation(messages=messages)
+        encoded = conversation.encode(
             tokenizer=self.tokenizer,
         )