diff --git a/memory_scope/models/llama_index_rerank_model.py b/memory_scope/models/llama_index_rank_model.py similarity index 97% rename from memory_scope/models/llama_index_rerank_model.py rename to memory_scope/models/llama_index_rank_model.py index c76cbd10..095fda5e 100644 --- a/memory_scope/models/llama_index_rerank_model.py +++ b/memory_scope/models/llama_index_rank_model.py @@ -9,7 +9,7 @@ from memory_scope.models.base_model import BaseModel, MODEL_REGISTRY from memory_scope.models.model_response import ModelResponse -class LlamaIndexRerankModel(BaseModel): +class LlamaIndexRankModel(BaseModel): m_type: ModelEnum = ModelEnum.RANK_MODEL MODEL_REGISTRY.register("dashscope_rank", DashScopeRerank) diff --git a/tests/models/test_models_lli_rerank.py b/tests/models/test_models_lli_rerank.py index c7bdf2f6..1238f7f8 100644 --- a/tests/models/test_models_lli_rerank.py +++ b/tests/models/test_models_lli_rerank.py @@ -1,7 +1,7 @@ import json import unittest -from memory_scope.models.llama_index_rerank_model import LlamaIndexRerankModel +from memory_scope.models.llama_index_rerank_model import LlamaIndexRankModel class TestLLIReRank(unittest.TestCase): @@ -13,7 +13,7 @@ class TestLLIReRank(unittest.TestCase): "model_name": "gte-rerank", "clazz": "models.llama_index_rerank_model" } - self.reranker = LlamaIndexRerankModel(**config) + self.reranker = LlamaIndexRankModel(**config) def test_rerank(self): query = "吃啥?"