Adds method to read the pooling types from model's files (#9506)
Signed-off-by: Flavia Beo <flavia.beo@ibm.com> Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Co-authored-by: Max de Bayser <mbayser@br.ibm.com>
This commit is contained in:
@@ -30,6 +30,13 @@ def test_limit_mm_per_prompt_parser(arg, expected):
|
||||
assert args.limit_mm_per_prompt == expected
|
||||
|
||||
|
||||
def test_valid_pooling_config():
|
||||
parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
|
||||
args = parser.parse_args(["--pooling-type=MEAN"])
|
||||
engine_args = EngineArgs.from_cli_args(args=args)
|
||||
assert engine_args.pooling_type == 'MEAN'
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("arg"),
|
||||
[
|
||||
|
||||
Reference in New Issue
Block a user