def optimize_model(self, model: FairseqModel) -> None:
gcfg = self.cfg.generation
model.make_generation_fast_(
beamable_mm_beam_size=None if gcfg.no_beamable_mm else gcfg.beam,
need_attn=gcfg.print_alignment,
)
if self.cfg.common.fp16:
model.half()
if not self.cfg.common.cpu:
model.cuda()
def __init__(self, dictionary: Dictionary, model: FairseqModel) -> None:
super().__init__()
self.dictionary = dictionary
self.model = model
self.unk = self.dictionary.unk()
self.save_incremental = False # this currently does not work properly
self.max_cache = 20_000
model.cuda()
model.eval()
model.make_generation_fast_()
self.states = {}
self.stateq = deque()
def optimize_model(self, model: FairseqModel) -> None:
model.make_generation_fast_()
if self.cfg.common.fp16:
model.half()
if not self.cfg.common.cpu:
model.cuda()