| | |
|---|
| | | |
|---|
| | | |
|---|
| | | from funasr.models.paraformer.search import Hypothesis |
|---|
| | | from funasr.models.sense_voice.utils.ctc_alignment import ctc_forced_align |
|---|
| | | from .utils.ctc_alignment import ctc_forced_align |
|---|
| | | |
|---|
| | | |
|---|
| | | class SinusoidalPositionEncoder(torch.nn.Module): |
|---|
| | |
|---|
| | | ): |
|---|
| | | """Embed positions in tensor.""" |
|---|
| | | maxlen = xs_pad.shape[1] |
|---|
| | | masks = sequence_mask(ilens, maxlen = maxlen, device=ilens.device)[:, None, :] |
|---|
| | | masks = sequence_mask(ilens, maxlen=maxlen, device=ilens.device)[:, None, :] |
|---|
| | | |
|---|
| | | xs_pad *= self.output_size() ** 0.5 |
|---|
| | | |
|---|
| | |
|---|
| | | |
|---|
| | | if output_timestamp: |
|---|
| | | from itertools import groupby |
|---|
| | | |
|---|
| | | timestamp = [] |
|---|
| | | tokens = tokenizer.text2tokens(text)[4:] |
|---|
| | | logits_speech = self.ctc.softmax(encoder_out)[i, 4:encoder_out_lens[i].item(), :] |
|---|
| | | token_back_to_id = tokenizer.tokens2ids(tokens) |
|---|
| | | token_ids = [] |
|---|
| | | for tok_ls in token_back_to_id: |
|---|
| | | if tok_ls: token_ids.extend(tok_ls) |
|---|
| | | else: token_ids.append(124) |
|---|
| | | |
|---|
| | | if len(token_ids) == 0: |
|---|
| | | result_i = {"key": key[i], "text": text} |
|---|
| | | results.append(result_i) |
|---|
| | | continue |
|---|
| | | |
|---|
| | | logits_speech = self.ctc.softmax(encoder_out)[i, 4 : encoder_out_lens[i].item(), :] |
|---|
| | | pred = logits_speech.argmax(-1).cpu() |
|---|
| | | logits_speech[pred==self.blank_id, self.blank_id] = 0 |
|---|
| | | logits_speech[pred == self.blank_id, self.blank_id] = 0 |
|---|
| | | align = ctc_forced_align( |
|---|
| | | logits_speech.unsqueeze(0).float(), |
|---|
| | | torch.Tensor(token_int[4:]).unsqueeze(0).long().to(logits_speech.device), |
|---|
| | | (encoder_out_lens-4).long(), |
|---|
| | | torch.tensor(len(token_int)-4).unsqueeze(0).long().to(logits_speech.device), |
|---|
| | | torch.Tensor(token_ids).unsqueeze(0).long().to(logits_speech.device), |
|---|
| | | (encoder_out_lens[i] - 4).long(), |
|---|
| | | torch.tensor(len(token_ids)).unsqueeze(0).long().to(logits_speech.device), |
|---|
| | | ignore_id=self.ignore_id, |
|---|
| | | ) |
|---|
| | | pred = groupby(align[0, :encoder_out_lens[0]]) |
|---|
| | | pred = groupby(align[0, : encoder_out_lens[i]]) |
|---|
| | | _start = 0 |
|---|
| | | token_id = 0 |
|---|
| | | ts_max = encoder_out_lens[i] - 4 |
|---|
| | | for pred_token, pred_frame in pred: |
|---|
| | | _end = _start + len(list(pred_frame)) |
|---|
| | | if pred_token != 0: |
|---|
| | | ts_left = max((_start*60-30)/1000, 0) |
|---|
| | | ts_right = min((_end*60-30)/1000, (ts_max*60-30)/1000) |
|---|
| | | ts_left = max((_start * 60 - 30) / 1000, 0) |
|---|
| | | ts_right = min((_end * 60 - 30) / 1000, (ts_max * 60 - 30) / 1000) |
|---|
| | | timestamp.append([tokens[token_id], ts_left, ts_right]) |
|---|
| | | token_id += 1 |
|---|
| | | _start = _end |
|---|
| | | timestamp = self.post(timestamp) |
|---|
| | | result_i = {"key": key[i], "text": text, "timestamp": timestamp} |
|---|
| | | timestamp, words = self.post(timestamp) |
|---|
| | | result_i = {"key": key[i], "text": text, "timestamp": timestamp, "words": words} |
|---|
| | | results.append(result_i) |
|---|
| | | else: |
|---|
| | | result_i = {"key": key[i], "text": text} |
|---|
| | |
|---|
| | | |
|---|
| | | def post(self, timestamp): |
|---|
| | | timestamp_new = [] |
|---|
| | | words_new = [] |
|---|
| | | prev_word = None |
|---|
| | | for i, t in enumerate(timestamp): |
|---|
| | | word, start, end = t |
|---|
| | | if word == '▁': |
|---|
| | | start = int(start * 1000) |
|---|
| | | end = int(end * 1000) |
|---|
| | | if word == "▁": |
|---|
| | | continue |
|---|
| | | if i == 0: |
|---|
| | | # timestamp_new.append([word, start, end]) |
|---|
| | | timestamp_new.append([int(start*1000), int(end*1000)]) |
|---|
| | | elif word.startswith("▁") or len(word) == 1 or not word[1].isalpha(): |
|---|
| | | timestamp_new.append([start, end]) |
|---|
| | | words_new.append(word) |
|---|
| | | elif word.startswith("▁"): |
|---|
| | | word = word[1:] |
|---|
| | | # timestamp_new.append([word, start, end]) |
|---|
| | | timestamp_new.append([int(start*1000), int(end*1000)]) |
|---|
| | | timestamp_new.append([start, end]) |
|---|
| | | words_new.append(word) |
|---|
| | | elif prev_word is not None and prev_word.isalpha() and prev_word.isascii() and word.isalpha() and word.isascii(): |
|---|
| | | word = prev_word + word |
|---|
| | | timestamp_new[-1][1] = end |
|---|
| | | words_new[-1] = word |
|---|
| | | else: |
|---|
| | | # timestamp_new[-1][0] += word |
|---|
| | | timestamp_new[-1][1] = int(end*1000) |
|---|
| | | return timestamp_new |
|---|
| | | timestamp_new.append([start, end]) |
|---|
| | | words_new.append(word) |
|---|
| | | prev_word = word |
|---|
| | | return timestamp_new, words_new |
|---|
| | | |
|---|
| | | def export(self, **kwargs): |
|---|
| | | from export_meta import export_rebuild_model |
|---|
| | | from .export_meta import export_rebuild_model |
|---|
| | | |
|---|
| | | if "max_seq_len" not in kwargs: |
|---|
| | | kwargs["max_seq_len"] = 512 |
|---|
| | |
|---|
| | | return models |
|---|
| | | |
|---|
| | | return results, meta_data |
|---|
| | | |
|---|
