| 662 | } |
| 663 | |
| 664 | static bool decode_helper(llama_context * ctx, llama_batch & batch, std::vector<float> & batch_logits, int n_batch, int n_vocab) { |
| 665 | int prev_outputs = 0; |
| 666 | for (int i = 0; i < (int) batch.n_tokens; i += n_batch) { |
| 667 | const int n_tokens = std::min<int>(n_batch, batch.n_tokens - i); |
| 668 | |
| 669 | llama_batch batch_view = { |
| 670 | n_tokens, |
| 671 | batch.token + i, |
| 672 | nullptr, |
| 673 | batch.pos + i, |
| 674 | batch.n_seq_id + i, |
| 675 | batch.seq_id + i, |
| 676 | batch.logits + i, |
| 677 | }; |
| 678 | |
| 679 | const int ret = llama_decode(ctx, batch_view); |
| 680 | if (ret != 0) { |
| 681 | LOG_ERR("failed to decode the batch, n_batch = %d, ret = %d\n", n_batch, ret); |
| 682 | return false; |
| 683 | } |
| 684 | |
| 685 | int n_outputs = 0; |
| 686 | for (int i = 0; i < n_tokens; ++i) { |
| 687 | n_outputs += batch_view.logits[i] != 0; |
| 688 | } |
| 689 | |
| 690 | memcpy(batch_logits.data() + size_t(prev_outputs)*n_vocab, llama_get_logits(ctx), size_t(n_outputs)*n_vocab*sizeof(float)); |
| 691 | |
| 692 | prev_outputs += n_outputs; |
| 693 | } |
| 694 | |
| 695 | return true; |
| 696 | } |
| 697 | |
| 698 | #define K_TOKEN_CHUNK 4 |
| 699 |
no test coverage detected