diff --git a/tests/test_qa_generator.py b/tests/test_qa_generator.py new file mode 100644 index 00000000..f715a674 --- /dev/null +++ b/tests/test_qa_generator.py @@ -0,0 +1,71 @@ +import importlib +import json +import sys +from types import ModuleType, SimpleNamespace + +import pytest +from pandas import Timestamp + +from weclone.data.models import Message, QaPair + + +def _load_data_processor(monkeypatch): + pii_module = ModuleType("weclone.core.PII.pii_detector") + pii_module.ChinesePIIDetector = object + pii_module.PIIDetector = object + monkeypatch.setitem(sys.modules, "weclone.core.PII.pii_detector", pii_module) + + cleaning_module = ModuleType("weclone.data.clean.strategies") + cleaning_module.LLMCleaningStrategy = object + cleaning_module.OlineLLMCleaningStrategy = object + monkeypatch.setitem(sys.modules, "weclone.data.clean.strategies", cleaning_module) + monkeypatch.delitem(sys.modules, "weclone.data.qa_generator", raising=False) + + return importlib.import_module("weclone.data.qa_generator").DataProcessor + + +def test_save_result_updates_configured_dataset_file(tmp_path, monkeypatch): + DataProcessor = _load_data_processor(monkeypatch) + dataset_dir = tmp_path / "custom-dataset" + dataset_dir.mkdir() + dataset_path = dataset_dir / "current-chat.json" + dataset_path.write_text('[{"messages": [{"content": "stale data"}]}]', encoding="utf-8") + (dataset_dir / "dataset_info.json").write_text( + json.dumps({"custom-chat": {"file_name": "current-chat.json"}}), + encoding="utf-8", + ) + + processor = DataProcessor.__new__(DataProcessor) + processor.c = SimpleNamespace(dataset="custom-chat", dataset_dir=str(dataset_dir)) + qa_pair = QaPair( + id=1, + time=Timestamp("2026-08-26T00:00:00"), + score=0, + messages=[Message(role="user", content="fresh data")], + images=[], + system="", + ) + monkeypatch.chdir(tmp_path) + + output_path = processor.save_result([qa_pair]) + + saved_data = json.loads(dataset_path.read_text(encoding="utf-8")) + assert saved_data[0]["messages"][0]["content"] == "fresh data" + assert output_path == str(dataset_path) + assert not (tmp_path / "dataset/res_csv/sft/sft-my.json").exists() + + +def test_save_result_rejects_non_object_dataset_entry(tmp_path, monkeypatch): + DataProcessor = _load_data_processor(monkeypatch) + dataset_dir = tmp_path / "custom-dataset" + dataset_dir.mkdir() + (dataset_dir / "dataset_info.json").write_text( + json.dumps({"custom-chat": None}), + encoding="utf-8", + ) + + processor = DataProcessor.__new__(DataProcessor) + processor.c = SimpleNamespace(dataset="custom-chat", dataset_dir=str(dataset_dir)) + + with pytest.raises(ValueError, match="must define file_name"): + processor.save_result([]) diff --git a/weclone/data/qa_generator.py b/weclone/data/qa_generator.py index e5a2d874..f9c409e0 100644 --- a/weclone/data/qa_generator.py +++ b/weclone/data/qa_generator.py @@ -146,11 +146,11 @@ def main(self): if self.enable_clean: self.clean_strategy.judge(qa_res) # type: ignore - self.save_result(qa_res) + output_path = self.save_result(qa_res) self._execute_length_cdf_script() logger.success( - f"Chat record processing successful, obtained {len(qa_res)} data entries in total, saved to ./dataset/res_csv/sft/sft-my.json" + f"Chat record processing successful, obtained {len(qa_res)} data entries in total, saved to {output_path}" ) def pre_parse_chat_dataset(self): @@ -678,7 +678,7 @@ def load_file(self, file_path) -> List[ChatMessage]: def process_text(self, chat_message: ChatMessage): pass - def save_result(self, qa_res: List[QaPair]): + def save_result(self, qa_res: List[QaPair]) -> str: """ Saves the list of QaPair objects to a JSON file after converting them to dictionaries. @@ -697,13 +697,23 @@ def save_result(self, qa_res: List[QaPair]): } processed_qa_res.append(item_dict) - output_path = "./dataset/res_csv/sft/sft-my.json" + dataset_info_path = os.path.join(self.c.dataset_dir, "dataset_info.json") + with open(dataset_info_path, "r", encoding="utf-8") as f: + dataset_info = json.load(f) + + dataset_entry = dataset_info.get(self.c.dataset) if isinstance(dataset_info, dict) else None + file_name = dataset_entry.get("file_name") if isinstance(dataset_entry, dict) else None + if not isinstance(file_name, str) or not file_name: + raise ValueError(f"Dataset '{self.c.dataset}' must define file_name in {dataset_info_path}") + + output_path = os.path.join(self.c.dataset_dir, file_name) os.makedirs(os.path.dirname(output_path), exist_ok=True) with open(output_path, "w", encoding="utf-8") as f: json.dump(processed_qa_res, f, ensure_ascii=False, indent=4) logger.success( f"Chat record processing successful, {len(qa_res)} entries in total, saved to {output_path}" ) + return output_path if __name__ == "__main__":