From 77b52e42cea4a31c0dfb605fbfa8c2a4ce27c79a Mon Sep 17 00:00:00 2001
From: enhuiz <niuzhe.nz@outlook.com>
Date: Mon, 16 Jan 2023 02:01:00 +0800
Subject: [PATCH] Batch accumulation. Ignore prompt and text loss in NAR and
 prompt loss in AR. Sampling temperature.

---
 vall_e/config.py      |  3 ++-
 vall_e/vall_e/ar.py   |  4 ++++
 vall_e/vall_e/base.py | 22 +++++++++++++++-------
 vall_e/vall_e/nar.py  |  4 ++++
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/vall_e/config.py b/vall_e/config.py
index 5e2b6c3..3ec6344 100644
--- a/vall_e/config.py
+++ b/vall_e/config.py
@@ -45,6 +45,7 @@ class Config(ConfigBase):
     max_phones: int = 50
 
     use_fp16: bool = True
+    gradient_accumulation_steps: int = 1
 
     @cached_property
     def get_spkr(self):
@@ -60,7 +61,7 @@ class Config(ConfigBase):
     def ds_cfg(self):
         return {
             "train_micro_batch_size_per_gpu": self.batch_size,
-            "gradient_accumulation_steps": 1,
+            "gradient_accumulation_steps": self.gradient_accumulation_steps,
             "optimizer": {
                 "type": "Adam",
                 "lr": self.warmup_min_lr,
diff --git a/vall_e/vall_e/ar.py b/vall_e/vall_e/ar.py
index 8f8c0fc..f06c15a 100644
--- a/vall_e/vall_e/ar.py
+++ b/vall_e/vall_e/ar.py
@@ -23,6 +23,10 @@ class AR(Base):
     def norm_type(self):
         return "ln"
 
+    @property
+    def resp_loss_only(self):
+        return False
+
     def _prune(self, l: Tensor):
         indices = (l == self.stop_token).nonzero()
         if len(indices) == 0:
diff --git a/vall_e/vall_e/base.py b/vall_e/vall_e/base.py
index f6df18c..a50c7cd 100644
--- a/vall_e/vall_e/base.py
+++ b/vall_e/vall_e/base.py
@@ -327,6 +327,10 @@ class Base(nn.Module):
     def n_prom_levels(self) -> int:
         return 8
 
+    @property
+    def resp_loss_only(self):
+        raise NotImplementedError
+
     def __init__(
         self,
         n_tokens: int,
@@ -334,7 +338,6 @@ class Base(nn.Module):
         n_heads: int = 8,
         n_layers: int = 12,
         p_dropout: float = 0.1,
-        resp_loss_only: bool = False,
     ):
         super().__init__()
         self.n_tokens = n_tokens
@@ -376,8 +379,6 @@ class Base(nn.Module):
 
         self.classifier = nn.Linear(d_model, n_resp_tokens)
 
-        self.resp_loss_only = resp_loss_only
-
     @property
     def stop_token(self):
         if not self.use_stop_token:
@@ -431,6 +432,7 @@ class Base(nn.Module):
         quant_levels: Tensor | None = None,
         shift_targ_list: bool = False,
         return_all_resp: bool = False,
+        sampling_temperature: float = 0.2,
     ):
         """
         Args:
@@ -441,6 +443,7 @@ class Base(nn.Module):
             quant_levels: specify which quant_levels to feed forward, used in NAR mode.
             shift_targ_list: whether to shift target list when computing loss. True if AR.
             return_all_resp: True if NAR.
+            sampling_temperature: a lower temperature makes the result more robust but less diverse.
         Returns:
             y: sampled tokens
         """
@@ -470,8 +473,11 @@ class Base(nn.Module):
 
             ignore_sep = torch.tensor(self.ignore_index, device=device)
 
-            # Predict the first level prom
-            prom_list = [t[..., 0] for t in proms_list]
+            # Ignore prom in the target
+            prom_list = [
+                torch.full_like(t[..., 0], self.ignore_index) for t in proms_list
+            ]
+
             text_prom_list = self._samplewise_merge_tensors(
                 text_list, prom_list, sep=ignore_sep
             )
@@ -506,9 +512,11 @@ class Base(nn.Module):
 
         if return_all_resp:
             logits = [hi[-li:] for hi, li in zip(h_list, map(len, resp_list))]
-            ret = [Categorical(logits=hi).sample() for hi in logits]
+            ret = [
+                Categorical(logits=hi / sampling_temperature).sample() for hi in logits
+            ]
         else:
             logits = torch.stack([hi[-1] for hi in h_list])
-            ret = Categorical(logits=logits).sample()
+            ret = Categorical(logits=logits / sampling_temperature).sample()
 
         return ret
diff --git a/vall_e/vall_e/nar.py b/vall_e/vall_e/nar.py
index 91ef8a9..4b303b1 100644
--- a/vall_e/vall_e/nar.py
+++ b/vall_e/vall_e/nar.py
@@ -22,6 +22,10 @@ class NAR(Base):
     def norm_type(self):
         return "adaln"
 
+    @property
+    def resp_loss_only(self):
+        return True
+
     def forward(
         self,
         text_list: list[Tensor],