From 2ddfc27d5b69e5c1a38021554e97dae958328c20 Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期四, 01 二月 2024 17:29:28 +0800
Subject: [PATCH] Funasr1.0 (#1343)
---
funasr/bin/train.py | 2 +-
funasr/models/whisper/model.py | 8 ++++++--
funasr/datasets/audio_datasets/samplers.py | 10 +++++++---
examples/industrial_data_pretraining/seaco_paraformer/demo.py | 3 ++-
4 files changed, 16 insertions(+), 7 deletions(-)
diff --git a/examples/industrial_data_pretraining/seaco_paraformer/demo.py b/examples/industrial_data_pretraining/seaco_paraformer/demo.py
index e9e226d..85d989e 100644
--- a/examples/industrial_data_pretraining/seaco_paraformer/demo.py
+++ b/examples/industrial_data_pretraining/seaco_paraformer/demo.py
@@ -25,6 +25,7 @@
# example2
import torchaudio
+import os
wav_file = os.path.join(model.model_path, "example/asr_example.wav")
input_tensor, sample_rate = torchaudio.load(wav_file)
input_tensor = input_tensor.mean(0)
@@ -33,7 +34,7 @@
# example3
import soundfile
-import os
+
wav_file = os.path.join(model.model_path, "example/asr_example.wav")
speech, sample_rate = soundfile.read(wav_file)
res = model.generate(input=[speech], batch_size_s=300, is_final=True)
diff --git a/funasr/bin/train.py b/funasr/bin/train.py
index d9d4d62..8ea0c0d 100644
--- a/funasr/bin/train.py
+++ b/funasr/bin/train.py
@@ -154,7 +154,7 @@
if batch_sampler is not None:
batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler)
batch_sampler = batch_sampler_class(dataset_tr, **kwargs.get("dataset_conf"))
- batch_sampler_val = batch_sampler_class(dataset_tr, is_training=False, **kwargs.get("dataset_conf"))
+ batch_sampler_val = batch_sampler_class(dataset_val, is_training=False, **kwargs.get("dataset_conf"))
dataloader_tr = torch.utils.data.DataLoader(dataset_tr,
collate_fn=dataset_tr.collator,
batch_sampler=batch_sampler,
diff --git a/funasr/datasets/audio_datasets/samplers.py b/funasr/datasets/audio_datasets/samplers.py
index 0d93098..535df5d 100644
--- a/funasr/datasets/audio_datasets/samplers.py
+++ b/funasr/datasets/audio_datasets/samplers.py
@@ -26,6 +26,8 @@
self.max_token_length = kwargs.get("max_token_length", 5000)
self.shuffle_idx = np.arange(self.total_samples)
self.shuffle = shuffle and is_training
+ self.length_scale_source = kwargs.get("length_scale_source", 1.0)
+
def __len__(self):
return (self.total_samples-1) // self.batch_size + 1
@@ -53,8 +55,10 @@
idx_map = self.shuffle_idx[idx]
# prompt = self.dataset.indexed_dataset[idx_map]["prompt"]
- sample_len_cur = self.dataset.get_source_len(idx_map) + \
- self.dataset.get_target_len(idx_map)
+ target_len = self.dataset.get_target_len(idx_map) if self.batch_type == 'length' else 0.0
+ source_len = self.dataset.get_source_len(idx_map) / self.length_scale_source
+ sample_len_cur = source_len + target_len
+
datalen_with_index.append([idx, sample_len_cur])
@@ -66,7 +70,7 @@
max_token_cur = max(max_token, sample_len_cur_raw)
max_token_padding = 1 + num_sample
- if self.batch_type == 'length':
+ if self.batch_type != 'example':
max_token_padding *= max_token_cur
if max_token_padding <= self.batch_size:
batch.append(idx)
diff --git a/funasr/models/whisper/model.py b/funasr/models/whisper/model.py
index 6a6d47c..381a501 100644
--- a/funasr/models/whisper/model.py
+++ b/funasr/models/whisper/model.py
@@ -10,6 +10,8 @@
from funasr.models.whisper.utils.decoding import detect_language as detect_language_function, decode as decode_function
+from funasr.register import tables
+
@dataclass
class ModelDimensions:
@@ -128,6 +130,8 @@
return x
+
+@tables.register("encoder_classes", "WhisperEncoder")
class AudioEncoder(nn.Module):
def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
super().__init__()
@@ -158,7 +162,7 @@
x = self.ln_post(x)
return x
-
+@tables.register("decoder_classes", "WhisperDecoder")
class TextDecoder(nn.Module):
def __init__(self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
super().__init__()
@@ -193,7 +197,7 @@
return logits
-
+@tables.register("model_classes", "Whisper")
class Whisper(nn.Module):
def __init__(self, dims: dict):
super().__init__()
--
Gitblit v1.9.1