for whatever reason this didn't work when I ran it on more than one at once... ugh.
(ctx, cache, probs)
| 767 | probs = context_output['probs'] |
| 768 | |
| 769 | def body(ctx, cache, probs): |
| 770 | """ for whatever reason this didn't work when I ran it on more than one at once... ugh.""" |
| 771 | next_outputs = sample_step(ctx[:, -1][:, None], ignore_ids=ignore_ids, news_config=news_config, |
| 772 | batch_size=batch_size, p_for_topp=p_for_topp, cache=cache, |
| 773 | do_topk=do_topk) |
| 774 | |
| 775 | |
| 776 | # Update everything |
| 777 | new_cache = tf.concat([cache, next_outputs['new_cache']], axis=-2) |
| 778 | new_ids = tf.concat([ctx, next_outputs['new_tokens'][:, None]], axis=1) |
| 779 | new_probs = tf.concat([probs, next_outputs['new_probs'][:, None]], axis=1) |
| 780 | return [new_ids, new_cache, new_probs] |
| 781 | |
| 782 | def cond(ctx, cache, probs): |
| 783 | # ctx = tf.Print(ctx,[tf.shape(ctx)]) |
nothing calls this directly
no test coverage detected