Spaces:
Sleeping
Sleeping
| from fam.llm.adapters.base import BaseDataAdapter | |
| class TiltedEncodec(BaseDataAdapter): | |
| def __init__(self, end_of_audio_token): | |
| self._end_of_audio_token = end_of_audio_token | |
| def decode(self, tokens: list[list[int]]) -> tuple[list[int], list[list[int]]]: | |
| assert len(tokens) > 1 | |
| text_ids = [] | |
| extracted_audio_ids = [] | |
| extracted_audio_ids.append([]) | |
| # Handle first hierarchy as special case as it contains text tokens as well | |
| # TODO: maybe it doesn't need special case, and can be handled on it's own :) | |
| for t in tokens[0]: | |
| if t > self._end_of_audio_token: | |
| text_ids.append(t) | |
| elif t < self._end_of_audio_token: | |
| extracted_audio_ids[0].append(t) | |
| # Handle the rest of the hierarchies | |
| for i in range(1, len(tokens)): | |
| token_hierarchy_ids = tokens[i] | |
| extracted_audio_ids.append([]) | |
| for t in token_hierarchy_ids: | |
| if t < self._end_of_audio_token: | |
| extracted_audio_ids[i].append(t) | |
| if len(set([len(x) for x in extracted_audio_ids])) != 1: | |
| min_len = min([len(x) for x in extracted_audio_ids]) | |
| max_len = max([len(x) for x in extracted_audio_ids]) | |
| print("WARNING: Number of tokens at each hierarchy must be of the same length!") | |
| print(f"Truncating to min length of {min_len} tokens from {max_len} max.") | |
| print([len(x) for x in extracted_audio_ids]) | |
| extracted_audio_ids = [x[:min_len] for x in extracted_audio_ids] | |
| return text_ids[:-1], extracted_audio_ids | |
| def encode(self, text_tokens: list[int], audio_tokens: list[list[int]]): | |
| """ | |
| Performs the required combination and padding as needed. | |
| """ | |
| raise NotImplementedError | |