mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-06-02 20:35:29 +00:00
fab9b931d9
* add infer_struct and infer_config * update codes * change InferConfig * Add hf_model_config to the engine * rm _get_hf_model_config * update codes * made adjustments according to the feedback from the reviewer. * update codes * add ci test for config and struct
38 lines
929 B
Python
38 lines
929 B
Python
from colossalai.inference.core.config import InferenceConfig
|
|
from colossalai.inference.core.inference_struct import BatchHandler, Sequence
|
|
|
|
|
|
def test_config_and_struct():
|
|
InferenceConfig("/llama")
|
|
sequence = Sequence(
|
|
request_id=1,
|
|
prompt="abc",
|
|
token_id=[1, 2, 3],
|
|
block_size=16,
|
|
sample_params=None,
|
|
block_table_index=1,
|
|
)
|
|
|
|
sequence2 = Sequence(
|
|
request_id=2,
|
|
prompt="bcd",
|
|
token_id=[4, 5, 6],
|
|
block_size=16,
|
|
sample_params=None,
|
|
block_table_index=2,
|
|
)
|
|
|
|
assert sequence.get_sentence_len() == 3
|
|
assert sequence.get_input_len() == 3
|
|
assert sequence.get_output_len() == 0
|
|
assert sequence.check_finish() == False
|
|
|
|
batch = BatchHandler.init_batch([sequence])
|
|
batch.fliter_batch()
|
|
batch.add_seqs([sequence2])
|
|
batch.clear_batch()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_config_and_struct()
|