diff --git a/config/demo_config.yaml b/config/demo_config.yaml index 1c45bce4..d36272ba 100644 --- a/config/demo_config.yaml +++ b/config/demo_config.yaml @@ -43,15 +43,27 @@ worker: generation_model: dashscope_generation embedding_model: dashscope_embedding rank_model: dashscope_rank + set_query_worker: + class: memory.worker.read.set_query_worker + extract_time_worker: + class: memory.worker.read.extract_time_worker + generation_model_top_k: 1 retrieve_store_worker: class: memory.worker.read.retrieve_store_worker retrieve_obs_top_k: 100 retrieve_ins_pf_top_k: 100 + semantic_rank_worker: + class: memory.worker.read.semantic_rank_worker fuse_rerank_worker: class: memory.worker.read.fuse_rerank_worker fuse_score_threshold: 0.1 fuse_ratio_dict: + conversation: 0.5 observation: 1 + obs_customized: 1 + insight: 2.0 + profile: 2.0 + profile_customized: 2.0 fuse_time_ratio: 2.0 fuse_rerank_top_k: 10 info_filter_worker: diff --git a/memory_scope/memory/worker/read/extract_time_worker.py b/memory_scope/memory/worker/read/extract_time_worker.py index f2368f11..4416eb93 100644 --- a/memory_scope/memory/worker/read/extract_time_worker.py +++ b/memory_scope/memory/worker/read/extract_time_worker.py @@ -30,7 +30,7 @@ class ExtractTimeWorker(MemoryBaseWorker): self.logger.info(f"extract_time_prompt={extract_time_prompt}") # call sft model - response = self.generation_model.call(prompt=extract_time_prompt, top_k=self.extra_time_top_k) + response = self.generation_model.call(prompt=extract_time_prompt, top_k=self.generation_model_top_k) # if empty, return if not response.status or not response.message.content: diff --git a/memory_scope/utils/prompt_handler.py b/memory_scope/utils/prompt_handler.py index 986022cc..521b6ed1 100644 --- a/memory_scope/utils/prompt_handler.py +++ b/memory_scope/utils/prompt_handler.py @@ -12,6 +12,7 @@ class PromptHandler(object): def __init__(self, class_path: str, prompt_file: str = "", prompt_dict: dict = None, **kwargs): self._class_path: str = class_path self._prompt_dict: Dict[str, str] = {} + self.kwargs = kwargs file_path = self._class_path.strip(".py") self.add_prompt_file(file_path) @@ -22,15 +23,31 @@ class PromptHandler(object): if prompt_dict: self.add_prompt_dict(prompt_dict) - def add_prompt_file(self, file_path: str): + @staticmethod + def file_path_completion(file_path: str) -> str: + if file_path.endswith(".yaml") or file_path.endswith(".json"): + return file_path + if os.path.exists(f"{file_path}.yaml"): - with open(f"{file_path}.yaml") as f: + return f"{file_path}.yaml" + + if os.path.exists(f"{file_path}.json"): + return f"{file_path}.json" + + raise RuntimeError(f"{file_path}/yaml/json is not exists!") + + def add_prompt_file(self, file_path: str): + file_path = self.file_path_completion(file_path) + + prompt_dict = {} + + if file_path.endswith(".yaml"): + with open(file_path) as f: prompt_dict = yaml.load(f, yaml.FullLoader) - elif os.path.exists(f"{file_path}.json"): + + elif file_path.endswith(".json"): with open(f"{file_path}.json") as f: prompt_dict = json.load(f) - else: - raise RuntimeError(f"{file_path}.yaml/json is not exists!") self.add_prompt_dict(prompt_dict) diff --git a/requirements.txt b/requirements.txt index 03159811..ee59e7f9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ -pyfiglet -termcolor +pyfiglet~=1.0.2 +termcolor~=2.4.0 llama-index==0.10.45 llama-index-core==0.10.44 llama-index-embeddings-dashscope==0.1.3 @@ -7,4 +7,9 @@ llama-index-llms-dashscope==0.1.2 llama-index-postprocessor-dashscope-rerank-custom==0.1.0 llama-index-vector-stores-elasticsearch==0.2.0 fire==0.6.0 -questionary==2.0.1 \ No newline at end of file +questionary==2.0.1 +requests~=2.31.0 +pydantic~=2.7.1 +dashscope~=1.19.1 +elasticsearch~=8.14.0 +pyyaml~=6.0.1 \ No newline at end of file