(
self,
input_ids: np.ndarray,
/,
*,
seq_id: int,
max_tokens: Optional[int],
)
| 1724 | return np.asarray(drafted, dtype=np.intc) |
| 1725 | |
| 1726 | def draft( |
| 1727 | self, |
| 1728 | input_ids: np.ndarray, |
| 1729 | /, |
| 1730 | *, |
| 1731 | seq_id: int, |
| 1732 | max_tokens: Optional[int], |
| 1733 | ) -> np.ndarray: |
| 1734 | if ( |
| 1735 | self.num_pred_tokens <= 0 |
| 1736 | or input_ids.size == 0 |
| 1737 | or seq_id < 0 |
| 1738 | or seq_id >= self.n_seq_max |
| 1739 | or not self.ready[seq_id] |
| 1740 | ): |
| 1741 | return np.array([], dtype=np.intc) |
| 1742 | n_predict = self.num_pred_tokens |
| 1743 | if max_tokens is not None: |
| 1744 | n_predict = min(n_predict, max_tokens) |
| 1745 | if n_predict <= 0: |
| 1746 | return np.array([], dtype=np.intc) |
| 1747 | |
| 1748 | n_past = int(input_ids.size) - 1 |
| 1749 | if self.ready_pos[seq_id] != n_past: |
| 1750 | return np.array([], dtype=np.intc) |
| 1751 | first_pos = n_past |
| 1752 | if first_pos < 0: |
| 1753 | return np.array([], dtype=np.intc) |
| 1754 | |
| 1755 | token = int(input_ids[-1]) |
| 1756 | drafted: List[int] = [] |
| 1757 | if self.chain_heads: |
| 1758 | return self._draft_chain_heads( |
| 1759 | seq_id=seq_id, |
| 1760 | first_pos=first_pos, |
| 1761 | token=token, |
| 1762 | n_predict=n_predict, |
| 1763 | ) |
| 1764 | if not self.is_mem_shared and self.context_pos[seq_id] > first_pos: |
| 1765 | self.truncate(seq_id, first_pos) |
| 1766 | if not self.is_mem_shared and self.context_pos[seq_id] < first_pos: |
| 1767 | self.ready[seq_id] = False |
| 1768 | return np.array([], dtype=np.intc) |
| 1769 | |
| 1770 | self._reset_sampler(seq_id) |
| 1771 | self._clear_batch() |
| 1772 | self._add_batch_token( |
| 1773 | token=token, |
| 1774 | pos=first_pos, |
| 1775 | seq_id=seq_id, |
| 1776 | logits=True, |
| 1777 | ) |
| 1778 | self._set_batch_embedding_row(0, self.pending_h[seq_id]) |
| 1779 | if not self._try_decode_batch(): |
| 1780 | if not self.is_mem_shared: |
| 1781 | self.truncate(seq_id, first_pos) |
| 1782 | return np.array([], dtype=np.intc) |
| 1783 | if not self.is_mem_shared: |
no test coverage detected