Merge pull request #157 from xming521/dev

xming521 · web-flow · commit 2defd5611edc · 2025-06-13T12:08:16.000+08:00
fix some bugs
diff --git a/tests/test_full_pipe.py b/tests/test_full_pipe.py
@@ -9,8 +9,8 @@
 
 import pytest
 
+from weclone.utils.config import load_config
 from weclone.utils.config_models import DataModality, WCMakeDatasetConfig
-from weclone.utils.configV2 import load_config
 from weclone.utils.log import logger
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
diff --git a/tests/test_full_pipeV2.py b/tests/test_full_pipeV2.py
@@ -9,8 +9,8 @@
 
 import pytest
 
+from weclone.utils.config import load_config
 from weclone.utils.config_models import DataModality, WCMakeDatasetConfig
-from weclone.utils.configV2 import load_config
 from weclone.utils.log import logger
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
diff --git a/weclone/cli.py b/weclone/cli.py
@@ -7,8 +7,8 @@
 import click
 import commentjson
 
+from weclone.utils.config import load_config
 from weclone.utils.config_models import CliArgs
-from weclone.utils.configV2 import load_config
 from weclone.utils.log import capture_output, logger
 
 cli_config: CliArgs | None = None
@@ -70,7 +70,7 @@ def cli():
 @apply_common_decorators()
 def qa_generator():
     """处理聊天记录CSV文件，生成问答对数据集。"""
-    from weclone.data.qa_generatorV2 import DataProcessor
+    from weclone.data.qa_generator import DataProcessor
 
     processor = DataProcessor()
     processor.main()
diff --git a/weclone/core/inference/offline_infer.py b/weclone/core/inference/offline_infer.py
@@ -42,22 +42,22 @@ def vllm_infer(
         raise ValueError("Pipeline parallel size should be smaller than the number of gpus.")
 
     model_args, data_args, _, generating_args = get_infer_args(
-        dict(
-            model_name_or_path=model_name_or_path,
-            adapter_name_or_path=adapter_name_or_path,
-            dataset=dataset,
-            dataset_dir=dataset_dir,
-            template=template,
-            cutoff_len=cutoff_len,
-            max_samples=max_samples,
-            preprocessing_num_workers=16,
-            vllm_config=vllm_config,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            max_new_tokens=max_new_tokens,
-            repetition_penalty=repetition_penalty,
-        )
+        {
+            "model_name_or_path": model_name_or_path,
+            "adapter_name_or_path": adapter_name_or_path,
+            "dataset": dataset,
+            "dataset_dir": dataset_dir,
+            "template": template,
+            "cutoff_len": cutoff_len,
+            "max_samples": max_samples,
+            "preprocessing_num_workers": 16,
+            "vllm_config": vllm_config,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "max_new_tokens": max_new_tokens,
+            "repetition_penalty": repetition_penalty,
+        }
     )
 
     tokenizer_module = load_tokenizer(model_args)
diff --git a/weclone/data/chat_parsers/wechat_parser.py b/weclone/data/chat_parsers/wechat_parser.py
@@ -7,7 +7,7 @@
 from tqdm import tqdm
 
 from weclone.data.models import ChatMessage
-from weclone.data.qa_generatorV2 import DataProcessor
+from weclone.data.qa_generator import DataProcessor
 from weclone.utils.log import logger
 
 data_dir = "./dataset/wechat/dat"
diff --git a/weclone/data/clean/strategies.py b/weclone/data/clean/strategies.py
@@ -113,8 +113,8 @@ def judge(self, data: List[QaPairV2]) -> None:
             inputs.append(prompt_value.to_string())
         outputs = vllm_infer(
             inputs,
-            self.make_dataset_config["model_name_or_path"],
-            template=self.make_dataset_config["template"],
+            self.make_dataset_config.model_name_or_path,
+            template=self.make_dataset_config.template,
             temperature=0,
             guided_decoding_class=QaPairScore,
             repetition_penalty=1.2,
diff --git a/weclone/data/qa_generator.py b/weclone/data/qa_generator.py
@@ -22,8 +22,8 @@
 )
 from weclone.data.strategies import LLMStrategy, TimeWindowStrategy
 from weclone.data.utils import ImageToTextProcessor, check_image_file_exists
+from weclone.utils.config import load_config
 from weclone.utils.config_models import DataModality, PlatformType, WCMakeDatasetConfig
-from weclone.utils.configV2 import load_config
 from weclone.utils.log import logger
 
 
diff --git a/weclone/eval/test_model.py b/weclone/eval/test_model.py
@@ -7,11 +7,12 @@
 from tqdm import tqdm
 
 from weclone.utils.config import load_config
+from weclone.utils.config_models import WCInferConfig
 
-config = load_config("web_demo")
+config = cast(WCInferConfig, load_config("web_demo"))
 
 config = {
-    "default_prompt": config["default_system"],
+    "default_prompt": config.default_system,
     "model": "gpt-3.5-turbo",
     "history_len": 15,
 }
diff --git a/weclone/eval/web_demo.py b/weclone/eval/web_demo.py
@@ -4,7 +4,7 @@
 
 
 def main():
-    config = load_config("web_demo")
+    load_config("web_demo")
     demo = create_web_demo()
     demo.queue()
     demo.launch(server_name="0.0.0.0", share=True, inbrowser=True)
diff --git a/weclone/train/train_sft.py b/weclone/train/train_sft.py
@@ -7,8 +7,8 @@
 from llamafactory.train.tuner import run_exp
 
 from weclone.data.clean.strategies import LLMCleaningStrategy
+from weclone.utils.config import load_config
 from weclone.utils.config_models import WCMakeDatasetConfig, WCTrainSftConfig
-from weclone.utils.configV2 import load_config
 from weclone.utils.log import logger
 
 
diff --git a/weclone/utils/config.py b/weclone/utils/config.py
@@ -1,67 +1,103 @@
 import os
 import sys
+from typing import Any, Dict, cast
 
 import commentjson
+from omegaconf import OmegaConf
+from pydantic import BaseModel
 
+from .config_models import (
+    WcConfig,
+    WCInferConfig,
+    WCMakeDatasetConfig,
+    WCTrainSftConfig,
+)
 from .log import logger
 from .tools import dict_to_argv
 
 
-def load_config(arg_type: str):
+def load_base_config() -> WcConfig:
+    """加载基础配置文件并创建WcConfig对象"""
     config_path = os.environ.get("WECLONE_CONFIG_PATH", "./settings.jsonc")
-    logger.info(f"Loading configuration from: {config_path}")  # Add logging to see which file is loaded
+    logger.info(f"Loading configuration from: {config_path}")
+
     try:
         with open(config_path, "r", encoding="utf-8") as f:
-            s_config: dict = commentjson.load(f)
+            s_config_dict: Dict[str, Any] = commentjson.load(f)
     except FileNotFoundError:
         logger.error(f"Configuration file not found: {config_path}")
-        sys.exit(1)  # Exit if config file is not found
+        sys.exit(1)
     except Exception as e:
         logger.error(f"Error loading configuration file {config_path}: {e}")
         sys.exit(1)
 
+    # 使用 OmegaConf 解析配置，然后转换为 Pydantic 模型验证
+    try:
+        omega_config = OmegaConf.create(s_config_dict)
+        config_dict_for_validation = OmegaConf.to_container(omega_config, resolve=True)
+        if not isinstance(config_dict_for_validation, dict):
+            raise TypeError(
+                f"Configuration should be a dictionary, but got {type(config_dict_for_validation)}"
+            )
+        wc_config = WcConfig(**cast(Dict[str, Any], config_dict_for_validation))
+    except Exception as e:
+        logger.error(f"Error parsing configuration with OmegaConf and WcConfig: {e}")
+        sys.exit(1)
+
+    return wc_config
+
+
+def create_config_by_arg_type(arg_type: str, wc_config: WcConfig) -> BaseModel:
+    """根据参数类型创建对应的配置对象,添加可能用到的参数,添加的参数会在model_validator中删除"""
     if arg_type == "cli_args":
-        config = s_config["cli_args"]
-    elif arg_type == "web_demo" or arg_type == "api_service":
-        # infer_args和common_args求并集
-        config = {**s_config["infer_args"], **s_config["common_args"]}
-    elif arg_type == "train_pt":
-        config = {**s_config["train_pt_args"], **s_config["common_args"]}
+        return wc_config.cli_args
+
+    common_config = wc_config.common_args.model_dump()
+
+    if arg_type == "web_demo" or arg_type == "api_service":
+        config_dict = {**common_config, **wc_config.infer_args.model_dump()}
+        return WCInferConfig(**config_dict)
+
     elif arg_type == "train_sft":
-        config = {**s_config["train_sft_args"], **s_config["common_args"]}
-        if s_config["make_dataset_args"]["prompt_with_history"]:
-            dataset_info_path = os.path.join(config["dataset_dir"], "dataset_info.json")
-            dataset_info = commentjson.load(open(dataset_info_path, "r", encoding="utf-8"))[config["dataset"]]
-            if dataset_info["columns"].get("history") is None:
-                logger.warning(
-                    f"{config['dataset']}数据集不包history字段，尝试使用wechat-sft-with-history数据集"
-                )
-                config["dataset"] = "wechat-sft-with-history"
-        if "image" in s_config["make_dataset_args"]["include_type"]:
-            if config["vision_api"].get("enable", False):
-                config["dataset"] = "wechat-img-rec-sft"  # 图像识别类模型使用的数据集
-            else:
-                config["dataset"] = "wechat-mllm-sft"  # 多模态模型使用的数据集
+        config_dict = {**common_config, **wc_config.train_sft_args.model_dump()}
+        return WCTrainSftConfig(**config_dict)
 
     elif arg_type == "make_dataset":
-        config = {**s_config["make_dataset_args"], **s_config["common_args"]}
-        config["dataset"] = s_config["train_sft_args"]["dataset"]
-        config["dataset_dir"] = s_config["train_sft_args"]["dataset_dir"]
-        config["cutoff_len"] = s_config["train_sft_args"]["cutoff_len"]
-        if "image" in config["include_type"]:
-            if config["vision_api"].get("enable", False):
-                config["dataset"] = "wechat-img-rec-sft"  # 图像识别类模型使用的数据集
-            else:
-                config["dataset"] = "wechat-mllm-sft"  # 多模态模型使用的数据集
+        make_dataset_config = wc_config.make_dataset_args.model_dump()
+        # ToDo 下面三个参数放到common里？
+        train_sft_args = wc_config.train_sft_args
+        extra_values = {
+            "dataset": train_sft_args.dataset,
+            "dataset_dir": train_sft_args.dataset_dir,
+            "cutoff_len": train_sft_args.cutoff_len,
+        }
+        config_dict = {**common_config, **make_dataset_config, **extra_values}
+        return WCMakeDatasetConfig(**config_dict)
 
     else:
         raise ValueError("暂不支持的参数类型")
 
-    if "train" in arg_type:
-        config["output_dir"] = config["adapter_name_or_path"]
-        config.pop("adapter_name_or_path")
-        config["do_train"] = True
 
-    sys.argv += dict_to_argv(config)
+def process_config_dict_and_argv(arg_type: str, config_pydantic: BaseModel) -> None:
+    """处理配置字典并更新sys.argv"""
+    config_dict = config_pydantic.model_dump(mode="json")
+
+    sys.argv += dict_to_argv(config_dict)
+
+
+def load_config(arg_type: str) -> BaseModel:
+    """加载配置的主函数"""
+    # 加载基础配置
+    wc_config = load_base_config()
+
+    # 根据类型创建配置对象
+    config_pydantic = create_config_by_arg_type(arg_type, wc_config)
+
+    # 处理配置字典和命令行参数
+    process_config_dict_and_argv(arg_type, config_pydantic)
+
+    return config_pydantic
+
 
-    return config
+if __name__ == "__main__":
+    load_config("train_sft")
diff --git a/weclone/utils/configV2.py b/weclone/utils/configV2.py
diff --git a/weclone/utils/config_models.py b/weclone/utils/config_models.py

Original file line number	Diff line number	Diff line change
`@@ -22,8 +22,8 @@`
`22`	`22`	`)`
`23`	`23`	`from weclone.data.strategies import LLMStrategy, TimeWindowStrategy`
`24`	`24`	`from weclone.data.utils import ImageToTextProcessor, check_image_file_exists`
	`25`	`+from weclone.utils.config import load_config`
`25`	`26`	`from weclone.utils.config_models import DataModality, PlatformType, WCMakeDatasetConfig`
`26`		`-from weclone.utils.configV2 import load_config`
`27`	`27`	`from weclone.utils.log import logger`
`28`	`28`
`29`	`29`