From 485abcb9663ea288465d96c6be2eed9e32df97e6 Mon Sep 17 00:00:00 2001 From: "xianzhe.xxz" Date: Tue, 18 Jun 2024 18:38:14 +0800 Subject: [PATCH] add llamaindex support --- memory_scope/db/elastic_search_client.py | 79 +++++++++++++++++- memory_scope/models/__init__.py | 18 ++++ memory_scope/models/dash_client.py | 33 ++++++++ memory_scope/models/dash_embedding_client.py | 62 +++++++++++++- memory_scope/models/dash_generate_client.py | 86 +++++++++++++++++++- memory_scope/models/dash_rerank_client.py | 78 +++++++++++++++++- memory_scope/utils/registry.py | 54 ++++++++++++ 7 files changed, 404 insertions(+), 6 deletions(-) create mode 100644 memory_scope/utils/registry.py diff --git a/memory_scope/db/elastic_search_client.py b/memory_scope/db/elastic_search_client.py index 74138cf3..f26116ce 100644 --- a/memory_scope/db/elastic_search_client.py +++ b/memory_scope/db/elastic_search_client.py @@ -1,11 +1,14 @@ from elasticsearch import Elasticsearch from elasticsearch.helpers import bulk -from common.dash_embedding_client import DashEmbeddingClient -from common.logger import Logger +from memory_scope.models.dash_embedding_client import DashEmbeddingClient, LLIEmbedding from constants.common_constants import ES_ENV_URL_DICT from enumeration.env_type import EnvType - +from memory_scope.utils.logger import Logger +from llama_index.core import VectorStoreIndex, StorageContext, ServiceContext +from llama_index.vector_stores.elasticsearch import ElasticsearchStore +from llama_index.core.schema import TextNode +from llama_index.vector_stores.elasticsearch import AsyncDenseVectorStrategy class ElasticSearchClient(object): def __init__(self, @@ -339,3 +342,73 @@ class ElasticSearchClient(object): self.print_hits(hits) return hits + + +class LLIElasticSearch(object): + def __init__(self, + es_index_name: str, + embedding_client: LLIEmbedding | None = None, + retrieve_topk: int = 3, + content_key: str = "text", + ): + self.es_index_name = es_index_name + self.content_key = content_key + self.embedding_client: LLIEmbedding = embedding_client + self.es_client = ElasticsearchStore(index_name="my_index", + es_url="http://localhost:9200", + retrieval_strategy=AsyncDenseVectorStrategy(hybrid=True)) + + self.service_context = ServiceContext.from_defaults(embed_model=self.embedding_client, llm=None) + self.storage_context = StorageContext.from_defaults(vector_store=self.es_client) + self.index = VectorStoreIndex(storage_context=self.storage_context, + service_context=self.service_context) + + self.retriever = self.index.as_retriever(similarity_top_k=retrieve_topk) + self.logger = Logger.get_logger() + + def log_index_info(self, ): + pass + + def print_hits(self, hits: list): + for hit in hits: + print_kwargs = { + "_id": hit['_id'], + "_score": hit['_score'], + } + for k, v in hit['_source'].items(): + # 不打印vector + if k == self.vector_key: + v = len(v) + print_kwargs[k] = v + self.logger.info(" ".join([f"{k}={v}" for k, v in print_kwargs.items()])) + + def similar_search(self, + text: str, + size: int, ): + + ret_nodes = self.retriever.retrieve(text) + return ret_nodes + + def insert_batch(self, doc_list:list[str]): + node_list = [] + for doc in doc_list: + assert "_id" in doc and "_source" in doc + content = doc["_source"]["text"] + doc["_source"].pop("text") + meta = doc["_source"] + node = TextNode(text=content, metadata=meta) + node.node_id(doc['_id']) + node_list.append(node) + self.index.insert_nodes(node_list) + + + def insert(self, _id: str, body: dict): + assert body and self.content_key in body, f"body={body} is illegal!" + content = body[self.content_key] + body.pop(self.content_key) + meta = body + node = TextNode(text=content, metadata=meta) + self.index.insert_nodes([node]) + + + \ No newline at end of file diff --git a/memory_scope/models/__init__.py b/memory_scope/models/__init__.py index e69de29b..73d7895f 100644 --- a/memory_scope/models/__init__.py +++ b/memory_scope/models/__init__.py @@ -0,0 +1,18 @@ +from utils.registry import Registry + +from llama_index.embeddings.dashscope import ( + DashScopeEmbedding, +) +# from llama_index.postprocessor.dashscope_rerank import DashScopeRerank + +from llama_index.llms.dashscope import DashScope # type: ignore + + +EMB = Registry('embedding') +EMB.register_module(DashScopeEmbedding) + +# RERANKER = Registry('reranker') +# RERANKER.register_module(DashScopeRerank) + +LLM = Registry('llm') +LLM.register_module(DashScope) diff --git a/memory_scope/models/dash_client.py b/memory_scope/models/dash_client.py index 6d07fed4..02206022 100644 --- a/memory_scope/models/dash_client.py +++ b/memory_scope/models/dash_client.py @@ -89,3 +89,36 @@ class DashClient(object): time.sleep(self.retry_sleep_time) return None + + +class LLIClient(object): + + def __init__(self, + model_name: str, + timeout: int = None, + max_retry_count: int = 2, + retry_sleep_time: float = 1.0, + **kwargs): + + self.model_name: str = model_name + self.timeout: int = timeout + self.max_retry_count: int = max_retry_count + self.retry_sleep_time: float = retry_sleep_time + self.kwargs: dict = kwargs + + self.data = {} + self.logger = Logger.get_logger() + + + def before_call(self, **kwargs): + pass + + def after_call(self, **kwargs): + pass + + def call_once(self, **kwargs): + pass + + def call(self, **kwargs): + pass + diff --git a/memory_scope/models/dash_embedding_client.py b/memory_scope/models/dash_embedding_client.py index 3529bfae..91ca82fd 100644 --- a/memory_scope/models/dash_embedding_client.py +++ b/memory_scope/models/dash_embedding_client.py @@ -1,8 +1,16 @@ from typing import List, Dict import dashscope +import time + +from models import EMB +from models.dash_client import DashClient, LLIClient + +from typing import List, Dict +from utils.registry import build_from_cfg +from utils.timer import Timer + -from common.dash_client import DashClient from constants.common_constants import DASH_ENV_URL_DICT, DASH_API_URL_DICT from enumeration.dash_api_enum import DashApiEnum @@ -41,3 +49,55 @@ class DashEmbeddingClient(DashClient): if len(embedding_results) == 1: embedding_results = list(embedding_results.values())[0] return embedding_results + +class LLIEmbedding(LLIClient): + + def __init__(self, method, model_name, **kwargs): + super(LLIEmbedding, self).__init__(model_name, **kwargs) + self.config = { + "method": method, + "model_name": model_name, + **kwargs} + self.embedder = build_from_cfg(self.config, EMB) + + 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, emb: Dict[int, List[float]], **kwargs) -> Dict[int, List[float]] | List[float]: + embedding_results = {} + for idx, e in enumerate(emb): + embedding_results[idx] = e + + if len(embedding_results) == 1: + embedding_results = list(embedding_results.values())[0] + return embedding_results + + + def call_once(self, model_name: str = None, retry_cnt: int = 0, **kwargs): + if model_name is None: + model_name = self.model_name + + self.before_call(model_name=model_name, **kwargs) + with Timer(self.__class__.__name__, log_time=False) as t: + self.logger.debug(f"data={self.data} timeout={self.timeout}") + try: + results = self.embedder.get_text_embedding_batch(**self.data) + results = self.after_call(results) + return results, True + except Exception as e: + self.logger.debug(f"Get Error in Embedding: {e}") + return None, False + + + def call(self, model_name: str = None, **kwargs): + for i in range(self.max_retry_count): + result, flag = self.call_once(model_name=model_name, retry_cnt=i, **kwargs) + if flag: + return result + else: + time.sleep(self.retry_sleep_time) + return None \ No newline at end of file diff --git a/memory_scope/models/dash_generate_client.py b/memory_scope/models/dash_generate_client.py index de6e11ec..6aefbe0f 100644 --- a/memory_scope/models/dash_generate_client.py +++ b/memory_scope/models/dash_generate_client.py @@ -2,10 +2,21 @@ from typing import List, Dict import dashscope -from common.dash_client import DashClient +from models.dash_client import DashClient, LLIClient from constants.common_constants import DASH_ENV_URL_DICT, DASH_API_URL_DICT from enumeration.dash_api_enum import DashApiEnum +import time +from typing import List, Dict +from utils.timer import Timer +from models import LLM +from utils.registry import build_from_cfg +from llama_index.core.base.llms.types import ChatMessage +from llama_index.core.base.llms.types import ( + ChatResponse, + CompletionResponse, +) + class DashGenerateClient(DashClient): """ @@ -43,3 +54,76 @@ class DashGenerateClient(DashClient): return output["choices"][0]["message"]["content"] else: raise NotImplementedError + + +class LLILLM(LLIClient): + + def __init__(self, method, model_name: str, **kwargs): + super(LLILLM, self).__init__(model_name, **kwargs) + self.config = { + "method": method, + "model_name": model_name, + **kwargs} + self.llm = build_from_cfg(self.config, LLM) + + + def before_call(self, model_name: str = None, **kwargs): + prompt: str = kwargs.pop("prompt", "") + messages: List[Dict[str, str]] = kwargs.pop("messages", []) + + if prompt: + input_text = prompt + input_type = 'prompt' + llama_input = input_text + elif messages: + input_text = messages + input_type = 'messages' + llama_input = [ChatMessage( + role=x['role'], content=x['content'] + ) for x in input_text] + else: + raise RuntimeError("prompt and messages is both empty!") + + self.data = { + input_type: llama_input, + } + + def after_call(self, response_obj, **kwargs): + self.logger.debug(f"response_obj={response_obj}") + if isinstance(response_obj, CompletionResponse): + return response_obj.text + elif isinstance(response_obj, ChatResponse): + return response_obj.message.content + else: + raise NotImplementedError + + + def call_once(self, model_name: str = None, retry_cnt: int = 0, **kwargs): + if model_name is None: + model_name = self.model_name + + self.before_call(model_name=model_name, **kwargs) + + with Timer(self.__class__.__name__, log_time=False) as t: + self.logger.debug(f"data={self.data} timeout={self.timeout}") + if True: + # try: + if 'prompt' in self.data: + results = self.llm.complete(**self.data) + else: + results = self.llm.chat(**self.data) + results = self.after_call(results) + return results, True + # except: + # return None, False + + + def call(self, model_name: str = None, **kwargs): + for i in range(self.max_retry_count): + result, flag = self.call_once(model_name=model_name, retry_cnt=i, **kwargs) + print("dashscope llm results:",result) + if flag: + return result + else: + time.sleep(self.retry_sleep_time) + return None \ No newline at end of file diff --git a/memory_scope/models/dash_rerank_client.py b/memory_scope/models/dash_rerank_client.py index be18766c..480d484f 100644 --- a/memory_scope/models/dash_rerank_client.py +++ b/memory_scope/models/dash_rerank_client.py @@ -2,11 +2,21 @@ from typing import List import dashscope -from common.dash_client import DashClient +from models.dash_client import DashClient, LLIClient from constants.common_constants import DASH_ENV_URL_DICT, DASH_API_URL_DICT from enumeration.dash_api_enum import DashApiEnum +import time +from typing import List +from models import RERANKER +from utils.timer import Timer + +from utils.registry import build_from_cfg +from llama_index.core.data_structs import Node +from llama_index.core.schema import NodeWithScore # type: ignore + + class DashReRankClient(DashClient): """ url: https://help.aliyun.com/document_detail/2780059.html @@ -41,3 +51,69 @@ class DashReRankClient(DashClient): def after_call(self, response_obj, **kwargs): return response_obj["output"]["results"] + + +class LLIReRank(LLIClient): + + def __init__(self, method, model_name, **kwargs): + super(LLIReRank, self).__init__(model_name, **kwargs) + + self.config = { + "method": method, + "model_name": model_name, + **kwargs} + self.reranker = build_from_cfg(self.config, RERANKER) + + + def before_call(self, model_name: str = None, **kwargs): + query: str = kwargs.pop("query", "") + documents: List[str] = kwargs.pop("documents", []) + top_n: int | None = kwargs.pop("top_n", None) + return_documents: bool = kwargs.pop("return_documents", False) + + assert query and documents, f"query or documents is empty! query={query}, documents={len(documents)}" + if top_n is None: + top_n = len(documents) + + nodes = [NodeWithScore(Node(text=text, score=1.0)) for text in documents] + self.reranker = self.reranker(top_n=top_n, + return_documents=return_documents) + + self.data = { + "nodes": nodes, + "query_str": query, + } + + + def after_call(self, nodes, **kwargs): + results = [] + for node in nodes: + results.append(dict(relevance_score=node.score, + document=node.node.text)) + return results + + + def call_once(self, model_name: str = None, retry_cnt: int = 0, **kwargs): + if model_name is None: + model_name = self.model_name + + self.before_call(model_name=model_name, **kwargs) + + with Timer(self.__class__.__name__, log_time=False) as t: + self.logger.debug(f"data={self.data} timeout={self.timeout}") + try: + results = self.reranker.postprocess_nodes(*self.data) + results = self.after_call(results) + return results, True + except: + return None, False + + + def call(self, model_name: str = None, **kwargs): + for i in range(self.max_retry_count): + result, flag = self.call_once(model_name=model_name, retry_cnt=i, **kwargs) + if flag: + return result + else: + time.sleep(self.retry_sleep_time) + return None \ No newline at end of file diff --git a/memory_scope/utils/registry.py b/memory_scope/utils/registry.py new file mode 100644 index 00000000..a662180f --- /dev/null +++ b/memory_scope/utils/registry.py @@ -0,0 +1,54 @@ +""" +Registry for different modules. +Init class according to the class name and verify the input parameters. +""" +import inspect + + +class Registry: + def __init__(self, name): + self.name = name + self.module_dict = dict() + + def register_module(self, module, module_name=None): + if module_name is None: + module_name = module.__name__ + if module_name in self.module_dict: + raise KeyError(f'{module_name} is already registered in {self.name}') + self.module_dict[module_name] = module + + def get_module(self, module_name): + assert module_name in self.module_dict, f'{module_name} not found in {self.name}' + return self.module_dict[module_name] + + +def build_from_cfg(config, registry, default_args: dict = None, skip_param_check=False): + + if default_args is None: + default_args = {} + + args = config.copy() + method_type = args.pop('method') + #params = args.get("parameters", {}) or default_args + params = args + if isinstance(method_type, str): + obj_cls = registry.get_module(method_type) + else: + raise TypeError( + f'type must be a str or valid type, but got {type(method_type)}') + + allowed_params = list(inspect.signature(obj_cls.__init__).parameters.keys()) + print(allowed_params, params) + if not skip_param_check: + filter_params = {key: value for key, value in params.items() if key in allowed_params} + else: + filter_params = params + # print( + # f"Registry {registry.name}, " + # f"allowed parameters {allowed_params}, filter parameters {filter_params}", + # flush=True + # ) + print(filter_params) + return obj_cls(**filter_params) + +