| from typing import List
|
|
|
| from tokenizers import AddedToken
|
| from transformers import WhisperTokenizer, WhisperProcessor
|
| import transformers.models.whisper.tokenization_whisper as whisper_tokenization
|
| from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE, TASK_IDS
|
|
|
| CUSTOM_TO_LANGUAGE_CODE = {**TO_LANGUAGE_CODE, "bambara": "bm"}
|
|
|
|
|
| whisper_tokenization.TO_LANGUAGE_CODE.update(CUSTOM_TO_LANGUAGE_CODE)
|
|
|
|
|
| class BambaraWhisperTokenizer(WhisperTokenizer):
|
| def __init__(self, *args, **kwargs):
|
| super().__init__(*args, **kwargs)
|
| self.add_tokens(AddedToken(content="<|bm|>", lstrip=False, rstrip=False, normalized=False, special=True))
|
|
|
| @property
|
| def prefix_tokens(self) -> List[int]:
|
| bos_token_id = self.convert_tokens_to_ids("<|startoftranscript|>")
|
| translate_token_id = self.convert_tokens_to_ids("<|translate|>")
|
| transcribe_token_id = self.convert_tokens_to_ids("<|transcribe|>")
|
| notimestamps_token_id = self.convert_tokens_to_ids("<|notimestamps|>")
|
|
|
| if self.language is not None:
|
| self.language = self.language.lower()
|
| if self.language in CUSTOM_TO_LANGUAGE_CODE:
|
| language_id = CUSTOM_TO_LANGUAGE_CODE[self.language]
|
| elif self.language in CUSTOM_TO_LANGUAGE_CODE.values():
|
| language_id = self.language
|
| else:
|
| is_language_code = len(self.language) == 2
|
| raise ValueError(
|
| f"Unsupported language: {self.language}. Language should be one of:"
|
| f" {list(CUSTOM_TO_LANGUAGE_CODE.values()) if is_language_code else list(CUSTOM_TO_LANGUAGE_CODE.keys())}."
|
| )
|
|
|
| if self.task is not None:
|
| if self.task not in TASK_IDS:
|
| raise ValueError(f"Unsupported task: {self.task}. Task should be in: {TASK_IDS}")
|
|
|
| bos_sequence = [bos_token_id]
|
| if self.language is not None:
|
| bos_sequence.append(self.convert_tokens_to_ids(f"<|{language_id}|>"))
|
| if self.task is not None:
|
| bos_sequence.append(transcribe_token_id if self.task == "transcribe" else translate_token_id)
|
| if not self.predict_timestamps:
|
| bos_sequence.append(notimestamps_token_id)
|
| return bos_sequence
|
|
|
