| | |
|---|
| | | class AutoModel: |
|---|
| | | |
|---|
| | | def __init__(self, **kwargs): |
|---|
| | | tables.print() |
|---|
| | | if kwargs.get("disable_log", False): |
|---|
| | | tables.print() |
|---|
| | | |
|---|
| | | model, kwargs = self.build_model(**kwargs) |
|---|
| | | |
|---|
| | |
|---|
| | | self.spk_model = spk_model |
|---|
| | | self.spk_kwargs = spk_kwargs |
|---|
| | | self.model_path = kwargs.get("model_path") |
|---|
| | | |
|---|
| | | |
|---|
| | | |
|---|
| | | def build_model(self, **kwargs): |
|---|
| | | assert "model" in kwargs |
|---|
| | |
|---|
| | | set_all_random_seed(kwargs.get("seed", 0)) |
|---|
| | | |
|---|
| | | device = kwargs.get("device", "cuda") |
|---|
| | | if not torch.cuda.is_available() or kwargs.get("ngpu", 0) == 0: |
|---|
| | | if not torch.cuda.is_available() or kwargs.get("ngpu", 1) == 0: |
|---|
| | | device = "cpu" |
|---|
| | | kwargs["batch_size"] = 1 |
|---|
| | | kwargs["device"] = device |
|---|
| | |
|---|
| | | res = self.model(*args, kwargs) |
|---|
| | | return res |
|---|
| | | |
|---|
| | | |
|---|
| | | |
|---|
| | | def generate(self, input, input_len=None, **cfg): |
|---|
| | | if self.vad_model is None: |
|---|
| | | return self.inference(input, input_len=input_len, **cfg) |
|---|
| | |
|---|
| | | data_batch = data_list[beg_idx:end_idx] |
|---|
| | | key_batch = key_list[beg_idx:end_idx] |
|---|
| | | batch = {"data_in": data_batch, "key": key_batch} |
|---|
| | | if (end_idx - beg_idx) == 1 and isinstance(data_batch[0], torch.Tensor): # fbank |
|---|
| | | if (end_idx - beg_idx) == 1 and kwargs.get("data_type", None) == "fbank": # fbank |
|---|
| | | batch["data_in"] = data_batch[0] |
|---|
| | | batch["data_lengths"] = input_len |
|---|
| | | |
|---|
