From 6b134ea703bc85ba4872e46c91aca69bbec75dc0 Mon Sep 17 00:00:00 2001
From: xming521 <1223398803@qq.com>
Date: Mon, 11 Aug 2025 11:03:14 +0800
Subject: [PATCH] feat(dataset): handle image data in LLM cleaning

Updates pre-commit hooks.

Previously, the cleaning pipeline was entirely disabled for datasets

containing image messages.

Image conversations now receive a default score, bypassing text-focused

LLM cleaning. This enables partial cleaning for mixed datasets and

prevents pipeline failures.
---
 .pre-commit-config.yaml          |  4 ++--
 weclone/data/clean/strategies.py | 19 +++++++++++--------
 weclone/data/qa_generator.py     |  4 ----
 3 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 9d360cf..58fb36e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -9,7 +9,7 @@ ci:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: check-ast # Python 语法检查
       - id: check-added-large-files # 防止大文件
@@ -27,7 +27,7 @@ repos:
         args: ["--fix=lf"]
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.12.7
+    rev: v0.12.8
     hooks:
       - id: ruff
         args: [--fix]
diff --git a/weclone/data/clean/strategies.py b/weclone/data/clean/strategies.py
index 06a6af5..ab952f7 100644
--- a/weclone/data/clean/strategies.py
+++ b/weclone/data/clean/strategies.py
@@ -86,14 +86,17 @@ class LLMCleaningStrategy(CleaningStrategy):
         inputs = []
         prompt_template = PromptTemplate.from_template(CLEAN_PROMPT)
         for qa in data:
-            messages_str = ""
-            for msg in qa.messages:
-                if msg.role == "user":
-                    messages_str += f"Q: {msg.content}\n"
-                elif msg.role == "assistant":
-                    messages_str += f"A: {msg.content}\n"
-            prompt_value = prompt_template.invoke({"id": qa.id, "messages": messages_str.strip()})
-            inputs.append(prompt_value.to_string())
+            if qa.images:
+                qa.score = 6
+            else:
+                messages_str = ""
+                for msg in qa.messages:
+                    if msg.role == "user":
+                        messages_str += f"Q: {msg.content}\n"
+                    elif msg.role == "assistant":
+                        messages_str += f"A: {msg.content}\n"
+                prompt_value = prompt_template.invoke({"id": qa.id, "messages": messages_str.strip()})
+                inputs.append(prompt_value.to_string())
 
         parsed_scores, failed_indexs = vllm_infer(
             inputs,
diff --git a/weclone/data/qa_generator.py b/weclone/data/qa_generator.py
index 8d930b6..6e3b4a0 100644
--- a/weclone/data/qa_generator.py
+++ b/weclone/data/qa_generator.py
@@ -89,10 +89,6 @@ class DataProcessor:
         clean_dataset_config = self.config.clean_dataset
 
         if self.enable_clean:
-            if DataModality.IMAGE in self.config.include_type:
-                logger.error("Enabling clean_dataset does not support image type messages")
-                exit()
-
             if clean_dataset_config.clean_strategy == "llm":
                 if self.config.online_llm_clear:
                     self.clean_strategy = OlineLLMCleaningStrategy(make_dataset_config=self.config)