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

Fix feeder

Browse Source
This commit is contained in:
enhuiz 2023-01-12 19:48:48 +08:00
parent 958c2df660
commit 777f9d6a13
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
vall_e/train.py
View File

@ -35,18 +35,20 @@ def main():
def train_feeder(engines, batch, name): def train_feeder(engines, batch, name):
model = engines["model"] model = engines["model"]
if cfg.model == "ar": if cfg.model.startswith("ar"):
_ = model( _ = model(
text_list=batch["text"], text_list=batch["text"],
proms_list=batch["proms"], proms_list=batch["proms"],
resp_list=batch["resp"], resp_list=batch["resp"],
) )
elif cfg.model == "nar": elif cfg.model.startswith("nar"):
_ = model( _ = model(
text_list=batch["text"], text_list=batch["text"],
proms_list=batch["proms"], proms_list=batch["proms"],
resps_list=batch["resps"], resps_list=batch["resps"],
) )
else:
raise NotImplementedError(cfg.model)
losses = model.gather_attribute("loss") losses = model.gather_attribute("loss")
@ -69,10 +71,10 @@ def main():
batch: dict batch: dict
batch = to_device(batch, cfg.device) batch = to_device(batch, cfg.device)
if cfg.model == "ar": if cfg.model.startswith("ar"):
resp_list = model(text_list=batch["text"], proms_list=batch["proms"]) resp_list = model(text_list=batch["text"], proms_list=batch["proms"])
resps_list = [r.unsqueeze(-1) for r in resp_list] resps_list = [r.unsqueeze(-1) for r in resp_list]
elif cfg.model == "nar": elif cfg.model.startswith("nar"):
resps_list = model( resps_list = model(
text_list=batch["text"], text_list=batch["text"],
proms_list=batch["proms"], proms_list=batch["proms"],