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 | |