mrq/vall-e
5
9
Fork 0
Code Issues 8 Pull Requests Packages Projects Releases Wiki Activity

when the

Browse Source
This commit is contained in:
mrq 2025-02-24 21:03:23 -06:00
parent 918e0dbac1
commit a5a04c39ef
1 changed files with 8 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
vall_e/models/base.py
View File

@ -1639,13 +1639,6 @@ class Base(nn.Module):
if loss_factor == 0.0:
continue
# cringe way to deduce "requested" level
level = quant_level
for i in range( self.n_resp_levels ):
if classifier_level == f'NAR:{i}:{i}':
level = i
break
if logits[batch_index].dim() < 3:
nll, metrics = _calc_loss( logits[batch_index][start:end], token.long(), causal )
@ -1653,6 +1646,12 @@ class Base(nn.Module):
if name == "resp":
name = f'{name}[{quant_level}]'
elif not self.resp_parallel_training:
# cringe way to deduce "requested" level
level = quant_level
for i in range( self.n_resp_levels ):
if classifier_level.endswith(f':{i}:{i}'):
level = i
break
if name == "resp":
name = f'{name}[{level}]'
sequence = token if token.dim() <= 1 else token[:, level]
@ -1711,9 +1710,10 @@ class Base(nn.Module):
# cringe way to deduce "requested" level
level = 0
for i in range( self.n_resp_levels ):
if classifier_level == f'NAR:{i}:{i}':
if classifier_level.endswith(f':{i}:{i}'):
level = i
break
sequence = [ x if x.dim() <= 1 else x[:, level] for x in target ]
sequence = _join( sequence, torch.tensor(self.ignore_index, device=sequence[-1].device) )
nll, metrics = _calc_loss( logits[batch_index][level], sequence.long(), causal )