10 lines
450 B
Python
10 lines
450 B
Python
|
import vllm.v1.sample.rejection_sampler as rs
|
||
|
|
|
||
|
|
from vllm_ascend.sample.rejection_sampler import apply_sampling_constraints, expand_batch_to_tokens, rejection_sample
|
||
|
|
|
||
|
|
# TODO: delete this patch after apply_sampling_constraints and rejection_sample
|
||
|
|
# are extracted to as class func of RejectionSampler
|
||
|
|
rs.apply_sampling_constraints = apply_sampling_constraints
|
||
|
|
rs.rejection_sample = rejection_sample
|
||
|
|
rs.expand_batch_to_tokens = expand_batch_to_tokens
|