o )i@s2ddlmZddlZddlZeGdddZdS)) dataclassNc@sxeZdZUejed<eeed<ejed<ejed<ejed<ejed<ddZe deeed ej d dfd d Z d S)SpecDecodeMetadatadraft_token_idsnum_draft_tokenscu_num_draft_tokenstarget_logits_indicesbonus_logits_indiceslogits_indicescCst|j|_dS)N)maxrZ max_spec_len)selfr h/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/vllm/v1/spec_decode/metadata.py __post_init__sz SpecDecodeMetadata.__post_init__devicereturnc Cst|}dd|D}t|g}t|}tj|tj|d}tj|tjd}t||} tj |tj|d} tj |tj|d} tj ||tj|d} |||| | | | dS)NcSsg|]}t|qSr )len).0idsr r r #sz1SpecDecodeMetadata.make_dummy..)dtyper)r)rrrrrr ) rsumtorchZtensorZint32npZcumsumZ from_numpytoZzeros) clsrrZ batch_sizerZflattened_draft_token_idsZ num_tokensZdraft_token_ids_tensorrZcu_num_draft_tokens_tensorrrr r r r make_dummys@   zSpecDecodeMetadata.make_dummyN) __name__ __module__ __qualname__rZTensor__annotations__listintr classmethodrrr r r r r s       r) dataclassesrnumpyrrrr r r r s