From 8fc58507678573abe6e87835b3579fa2b9dc39a2 Mon Sep 17 00:00:00 2001 From: Shengsheng Huang Date: Wed, 24 Apr 2024 21:41:41 +0800 Subject: [PATCH] fix test format --- .../tests/integration_tests/llms/test_bigdl_llm.py | 2 +- .../tests/integration_tests/llms/test_ipex_llm.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/libs/community/tests/integration_tests/llms/test_bigdl_llm.py b/libs/community/tests/integration_tests/llms/test_bigdl_llm.py index 4162fbdb89ed9..e40f5488054c3 100644 --- a/libs/community/tests/integration_tests/llms/test_bigdl_llm.py +++ b/libs/community/tests/integration_tests/llms/test_bigdl_llm.py @@ -11,7 +11,7 @@ not model_ids_to_test, reason="TEST_BIGDLLLM_MODEL_IDS environment variable not set.", ) -model_ids_to_test = [model_id.strip() for model_id in model_ids_to_test.split(",")] +model_ids_to_test = [model_id.strip() for model_id in model_ids_to_test.split(",")] # type: ignore @skip_if_no_model_ids diff --git a/libs/community/tests/integration_tests/llms/test_ipex_llm.py b/libs/community/tests/integration_tests/llms/test_ipex_llm.py index c895be9f60ce1..c57c5fba74b74 100644 --- a/libs/community/tests/integration_tests/llms/test_ipex_llm.py +++ b/libs/community/tests/integration_tests/llms/test_ipex_llm.py @@ -1,5 +1,6 @@ """Test IPEX LLM""" import os +from typing import Any import pytest from langchain_core.outputs import LLMResult @@ -10,10 +11,10 @@ skip_if_no_model_ids = pytest.mark.skipif( not model_ids_to_test, reason="TEST_IPEXLLM_MODEL_IDS environment variable not set." ) -model_ids_to_test = [model_id.strip() for model_id in model_ids_to_test.split(",")] +model_ids_to_test = [model_id.strip() for model_id in model_ids_to_test.split(",")] # type: ignore -def load_model(model_id: str) -> None: +def load_model(model_id: str) -> Any: llm = IpexLLM.from_model_id( model_id=model_id, model_kwargs={"temperature": 0, "max_length": 16, "trust_remote_code": True}, @@ -21,7 +22,7 @@ def load_model(model_id: str) -> None: return llm -def load_model_more_types(model_id: str, load_in_low_bit: str) -> None: +def load_model_more_types(model_id: str, load_in_low_bit: str) -> Any: llm = IpexLLM.from_model_id( model_id=model_id, load_in_low_bit=load_in_low_bit,