diff --git a/memory_scope/models/base_embedding_model.py b/memory_scope/models/base_embedding_model.py index 3e6d30cf..6144e2a7 100644 --- a/memory_scope/models/base_embedding_model.py +++ b/memory_scope/models/base_embedding_model.py @@ -1,3 +1,4 @@ +from typing import List, Dict from llama_index.embeddings.dashscope import DashScopeEmbedding from memory_scope.models import MODEL_REGISTRY @@ -22,3 +23,31 @@ class BaseEmbeddingModel(BaseModel): async def _async_call(self, **kwargs) -> ModelResponse: pass + + +class LLIEmbedding(BaseEmbeddingModel): + + def before_call(self, **kwargs): + text: str | List[str] = kwargs.pop("text", "") + if isinstance(text, str): + text = [text] + self.data = dict(texts=text) + + def after_call(self, model_response: ModelResponse, **kwargs) -> ModelResponse: + embeddings = model_response.raw + model_response.embedding_results = embeddings + return model_response + + def _call(self, **kwargs) -> ModelResponse: + results = ModelResponse() + try: + response = self.model.get_text_embedding_batch(**self.data) + results.raw = response + results.status = True + except Exception as e: + results.details = e + results.status = False + return results + + + \ No newline at end of file diff --git a/memory_scope/models/base_generation_model.py b/memory_scope/models/base_generation_model.py index 0639f29a..a8881cba 100644 --- a/memory_scope/models/base_generation_model.py +++ b/memory_scope/models/base_generation_model.py @@ -31,7 +31,7 @@ class BaseGenerationModel(BaseModel): pass -class LLILLM(BaseGenerateModel): +class LLILLM(BaseGenerationModel): def before_call(self, **kwargs) -> None: prompt: str = kwargs.pop("prompt", "") diff --git a/memory_scope/models/base_rank_model.py b/memory_scope/models/base_rank_model.py index ce6326cd..c9ed0dfe 100644 --- a/memory_scope/models/base_rank_model.py +++ b/memory_scope/models/base_rank_model.py @@ -1,3 +1,7 @@ +from typing import List, Dict +from llama_index.core.data_structs import Node +from llama_index.core.schema import NodeWithScore # type: ignore + from memory_scope.models import MODEL_REGISTRY from memory_scope.models.base_model import BaseModel from memory_scope.models.response import ModelResponse, ModelResponseGen @@ -11,12 +15,52 @@ class BaseRankModel(BaseModel): def before_call(self, **kwargs) -> None: pass - def after_call(self, model_response: ModelResponse | ModelResponseGen, - **kwargs) -> ModelResponse | ModelResponseGen: + def after_call(self, **kwargs) -> ModelResponse: pass - def _call(self, stream: bool = False, **kwargs) -> ModelResponse | ModelResponseGen: + def _call(self, stream: bool = False, **kwargs) -> ModelResponse: pass async def _async_call(self, **kwargs) -> ModelResponse: pass + + +class LLIReRank(BaseRankModel): + + def before_call(self, **kwargs) -> None: + assert "query" in kwargs or "documents" in kwargs + query: str = kwargs.pop("query", "") + documents: List[str] = kwargs.pop("documents", []) + + assert query and documents, f"query or documents is empty! query={query}, documents={len(documents)}" + if top_n is None: + top_n = len(documents) + + # using -1.0 as dummy scores + nodes = [NodeWithScore(node=Node(text=text), score=-1.0) for text in documents] + + self.data = { + "nodes": nodes, + "query_str": query, + } + + def after_call(self, nodes: List[NodeWithScore]) -> ModelResponse: + ranks = list() + for node in nodes: + ranks.append(dict(relevance_score=node.score, + document=node.node.text)) + results = ModelResponse(rank_scores=ranks) + return results + + def _call(self, **kwargs) -> ModelResponse: + results = ModelResponse() + try: + self.before_call(**kwargs) + response = self.model.postprocess_nodes(**self.data) + response = self.after_call(response) + results.rank_scores = response + results.status = True + except Exception as e: + results.details = e + results.status = False + return results diff --git a/memory_scope/models/response.py b/memory_scope/models/response.py index 5c435042..52016c89 100644 --- a/memory_scope/models/response.py +++ b/memory_scope/models/response.py @@ -1,4 +1,4 @@ -from typing import Generator, List, Dict +from typing import Generator, List, Dict, Any from pydantic import BaseModel, Field @@ -6,10 +6,11 @@ from pydantic import BaseModel, Field class ModelResponse(BaseModel): text: str = Field("", description="") - embedding_results: Dict[int, List[float]] | List[float] = Field([], description="") - - rank_scores: Dict[int, float] = Field({}, description="The rank scores of each documents.") + embedding_results: List[List[float]] = Field([], description="") + #rank_scores: Dict[int, float] = Field({}, description="The rank scores of each documents.") + rank_scores: List[Dict[str, Any]] = Field({}, description="The rank scores of each documents.") + # [{"document": "xxx", "score": 0.5}, {{"document": "yyy", "score": 0.3}}] model_type: str = Field("", description="One of LLM, EMB, RANK.") status: bool = Field(True, description="Indicates whether the model call was successful.") @@ -17,5 +18,5 @@ class ModelResponse(BaseModel): details: str = Field("", description=("The details information for model call, \ usually for storage of raw response or failure messages.")) - + raw: Any = Field("", description=("raw response from model call")) ModelResponseGen = Generator[ModelResponse, None, None] diff --git a/memory_scope/utils/registry.py b/memory_scope/utils/registry.py index 0eb8ee75..8c9d14df 100644 --- a/memory_scope/utils/registry.py +++ b/memory_scope/utils/registry.py @@ -22,6 +22,6 @@ class Registry(object): module_name_dict = {m.__name__: m for m in modules} self.module_dict.update(module_name_dict) - def __getitem__(self, module_name: str): + def get(self, module_name: str): assert module_name in self.module_dict, f'{module_name} not found in {self.name}' return self.module_dict[module_name] diff --git a/tests/models/test_models_llm_embedding.py b/tests/models/test_models_llm_embedding.py new file mode 100644 index 00000000..9d08bee9 --- /dev/null +++ b/tests/models/test_models_llm_embedding.py @@ -0,0 +1,24 @@ +import unittest +from memory_scope.models.base_embedding_model import LLIEmbedding + +class TestLLIEmbedding(unittest.TestCase): + """Tests for LLIEmbedding""" + + def setUp(self): + config = { + "method_type": "DashScopeEmbedding", + "model_name": "text-embedding-v2", + "clazz": "models.base_embedding_model" + } + self.emb = LLIEmbedding(**config) + + def test_single_embedding(self): + text = "您吃了吗？" + embs = self.emb.call(text=text) + + def test_batch_embedding(self): + text = ["您吃了吗？", + "吃了吗您？"] + embs = self.emb.call(text=text) + + \ No newline at end of file