(
self,
requests: Sequence[Tuple[np.ndarray, int, Optional[int]]],
/,
)
| 1816 | return np.asarray(drafted, dtype=np.intc) |
| 1817 | |
| 1818 | def draft_many( |
| 1819 | self, |
| 1820 | requests: Sequence[Tuple[np.ndarray, int, Optional[int]]], |
| 1821 | /, |
| 1822 | ) -> List[np.ndarray]: |
| 1823 | results = [np.array([], dtype=np.intc) for _ in requests] |
| 1824 | if self.chain_heads: |
| 1825 | for result_index, (input_ids, seq_id, max_tokens) in enumerate(requests): |
| 1826 | results[result_index] = self.draft( |
| 1827 | input_ids, |
| 1828 | seq_id=seq_id, |
| 1829 | max_tokens=max_tokens, |
| 1830 | ) |
| 1831 | return results |
| 1832 | |
| 1833 | active: List["MTPDraftProvider.DraftManyState"] = [] |
| 1834 | for result_index, (input_ids, seq_id, max_tokens) in enumerate(requests): |
| 1835 | if ( |
| 1836 | self.num_pred_tokens <= 0 |
| 1837 | or input_ids.size == 0 |
| 1838 | or seq_id < 0 |
| 1839 | or seq_id >= self.n_seq_max |
| 1840 | or not self.ready[seq_id] |
| 1841 | ): |
| 1842 | continue |
| 1843 | n_predict = self.num_pred_tokens |
| 1844 | if max_tokens is not None: |
| 1845 | n_predict = min(n_predict, max_tokens) |
| 1846 | if n_predict <= 0: |
| 1847 | continue |
| 1848 | if len(active) >= self.n_batch: |
| 1849 | break |
| 1850 | |
| 1851 | n_past = int(input_ids.size) - 1 |
| 1852 | if self.ready_pos[seq_id] != n_past: |
| 1853 | continue |
| 1854 | first_pos = n_past |
| 1855 | if first_pos < 0: |
| 1856 | continue |
| 1857 | if not self.is_mem_shared and self.context_pos[seq_id] > first_pos: |
| 1858 | self.truncate(seq_id, first_pos) |
| 1859 | if not self.is_mem_shared and self.context_pos[seq_id] < first_pos: |
| 1860 | self.ready[seq_id] = False |
| 1861 | continue |
| 1862 | self._reset_sampler(seq_id) |
| 1863 | active.append( |
| 1864 | self.DraftManyState( |
| 1865 | result_index=result_index, |
| 1866 | seq_id=seq_id, |
| 1867 | first_pos=first_pos, |
| 1868 | keep_len=n_past, |
| 1869 | n_predict=n_predict, |
| 1870 | token=int(input_ids[-1]), |
| 1871 | drafted=[], |
| 1872 | embedding=self.pending_h[seq_id], |
| 1873 | cache_key=( |
| 1874 | tuple(int(token) for token in input_ids.tolist()), |
| 1875 | n_predict, |
nothing calls this directly
no test coverage detected