THUDM / SwissArmyTransformer

SwissArmyTransformer is a flexible and powerful library to develop your own Transformer variants.
https://THUDM.github.io/SwissArmyTransformer
Apache License 2.0
987 stars 95 forks source link

stream_filling_sequence function #114

Closed wenyihong closed 1 year ago

wenyihong commented 1 year ago

Usage shown in func stream_filling_sequence's comments

filling_stream = stream_filling_sequence(
                model, seq,
                batch_size=1,
                get_masks_and_position_ids=get_func,
                strategy=strategy,
                pre_image=pre_image,
                image=torch_image,
            )
for tokens, logit, mems in filling_stream:
     pass # you can do something or nothing at all
output = strategy.finalize(tokens, mems)[0]