WAC-GEC

app.py

@@ -11,6 +11,17 @@ import spacy
 from spellchecker import SpellChecker
 import difflib
 
+# ======================== 新增：WAC-GEC导入 ========================
+try:
+    from whitespace_correction import WhitespaceCorrector
+    WAC_GEC_AVAILABLE = True
+    # 初始化WAC-GEC模型（使用CPU，HF Space通常没有GPU）
+    wac_corrector = None  # 延迟初始化
+except ImportError:
+    WAC_GEC_AVAILABLE = False
+    wac_corrector = None
+    print("⚠️ whitespace_correction未安装，WAC-GEC功能将不可用")
+
 # ======================== API配置 ========================
 DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY", "")
 DEEPSEEK_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
@@ -60,6 +71,56 @@ Next, please correct the following sentence according to the above requirements.
 
 [input]: """
 
+# ======================== 新增：WAC-GEC初始化函数 ========================
+def initialize_wac_gec():
+    """延迟初始化WAC-GEC模型"""
+    global wac_corrector
+    if not WAC_GEC_AVAILABLE:
+        return False
+    
+    if wac_corrector is None:
+        try:
+            # 根据环境选择设备
+            device = "cpu"  # HF Space默认使用CPU
+            # 如果有GPU可用，取消下面的注释
+            # import torch
+            # device = "cuda:0" if torch.cuda.is_available() else "cpu"
+            
+            # 优先使用本地模型，如果不存在则自动下载
+            local_model_path = "./models"  # HF Space中的模型目录
+            if os.path.exists(os.path.join(local_model_path, "eo_larger_byte")):
+                wac_corrector = WhitespaceCorrector.from_pretrained(
+                    model="eo_larger_byte",
+                    device=device,
+                    download_dir=local_model_path
+                )
+            else:
+                # 如果本地没有，自动下载到默认缓存
+                wac_corrector = WhitespaceCorrector.from_pretrained(
+                    model="eo_larger_byte",
+                    device=device,
+                    download_dir=None
+                )
+            print(f"✅ WAC-GEC模型已加载 (设备: {device})")
+            return True
+        except Exception as e:
+            print(f"❌ WAC-GEC模型加载失败: {e}")
+            return False
+    return True
+
+# ======================== 新增：WAC-GEC处理函数 ========================
+def call_wac_gec(text):
+    """使用WAC-GEC纠正空白符错误"""
+    if not initialize_wac_gec():
+        raise ValueError("⚠️ WAC-GEC模型未安装或加载失败")
+    
+    try:
+        corrected = wac_corrector.correct_text(text)
+        # 格式化输出以匹配DeepSeek的格式
+        return f"[output]: {corrected}"
+    except Exception as e:
+        raise Exception(f"WAC-GEC处理错误: {str(e)}")
+
 # ======================== 新增：颜色对比函数 ========================
 def generate_colored_diff(original, cleaned):
     """
@@ -163,12 +224,13 @@ def create_comparison_html(original_list, cleaned_list):
     return html
 
 # ======================== 工具函数 ========================
-def check_api_key():
-    if not DEEPSEEK_API_KEY:
+def check_api_key(model_choice):
+    """检查API密钥（仅DeepSeek需要）"""
+    if model_choice == "deepseek-r1-distill-llama-8b" and not DEEPSEEK_API_KEY:
         raise ValueError("⚠️ 请在 Space Settings 中配置 DEEPSEEK_API_KEY！")
 
 def call_deepseek_api(prompt, model="deepseek-r1-distill-llama-8b", temperature=0.1, stream=True):
-    check_api_key()
+    check_api_key(model)
     client = OpenAI(api_key=DEEPSEEK_API_KEY, base_url=DEEPSEEK_BASE_URL)
     completion = client.chat.completions.create(
         model=model,
@@ -286,13 +348,19 @@ def search_leaderboard(df, query):
         return df
     return df[df['Benchmark'].str.contains(query, case=False, na=False)]
 
-# ======================== 数据清洗函数（修改版）========================
+# ======================== 数据清洗函数（修改版：支持双模型）========================
 def clean_dataset(file_path, question_column, model_choice, temperature, max_samples, progress=gr.Progress()):
     try:
+        # 检查API密钥（仅DeepSeek需要）
         try:
-            check_api_key()
+            check_api_key(model_choice)
         except ValueError as e:
-            return str(e), None, ""
+            if model_choice == "deepseek-r1-distill-llama-8b":
+                return str(e), None, ""
+        
+        # 检查WAC-GEC可用性
+        if model_choice == "WAC-GEC" and not WAC_GEC_AVAILABLE:
+            return "❌ WAC-GEC模型未安装！请安装 whitespace_correction 包。", None, ""
         
         progress(0.05, desc="📁 读取数据文件...")
         df = pd.read_parquet(file_path)
@@ -309,12 +377,18 @@ def clean_dataset(file_path, question_column, model_choice, temperature, max_sam
         war_original = calculate_whitespace_anomaly_rate(original_sentences)
         sed_original = calculate_spelling_error_density(original_sentences)
         
-        progress(0.1, desc=f"🚀 开始清洗 {total} 个样本...")
+        progress(0.1, desc=f"🚀 开始清洗 {total} 个样本 (模型: {model_choice})...")
+        
+        # WAC-GEC不需要添加___标记
+        if model_choice == "WAC-GEC":
+            data_corrupt = [str(item) for item in data_ori]
+        else:
+            data_corrupt = [process_sentence(str(item)) for item in data_ori]
         
-        data_corrupt = [process_sentence(str(item)) for item in data_ori]
         results = []
-        max_retries = 5
-        log_text = f"🚀 开始处理 {total} 个样本...\n\n"
+        max_retries = 5 if model_choice == "deepseek-r1-distill-llama-8b" else 3
+        log_text = f"🚀 开始处理 {total} 个样本...\n"
+        log_text += f"📌 使用模型: {model_choice}\n\n"
         
         for idx in range(total):
             progress((0.1 + 0.7 * idx / total), desc=f"处理中: {idx+1}/{total}")
@@ -326,21 +400,33 @@ def clean_dataset(file_path, question_column, model_choice, temperature, max_sam
             
             while retry_count < max_retries:
                 try:
-                    response_content = call_deepseek_api(
-                        PROMPT_TEMPLATE + original_text,
-                        model=model_choice,
-                        temperature=float(temperature)
-                    )
+                    # 根据模型选择调用不同的API
+                    if model_choice == "WAC-GEC":
+                        response_content = call_wac_gec(original_text)
+                    else:
+                        response_content = call_deepseek_api(
+                            PROMPT_TEMPLATE + original_text,
+                            model=model_choice,
+                            temperature=float(temperature)
+                        )
                     
-                    if is_valid_output(response_content, original_text, unprocess_text):
-                        results.append(response_content)
-                        break
+                    # WAC-GEC的输出格式简单，无需复杂验证
+                    if model_choice == "WAC-GEC":
+                        if response_content.startswith('[output]:'):
+                            results.append(response_content)
+                            break
+                        else:
+                            retry_count += 1
                     else:
-                        retry_count += 1
+                        if is_valid_output(response_content, original_text, unprocess_text):
+                            results.append(response_content)
+                            break
+                        else:
+                            retry_count += 1
                         
                 except Exception as e:
                     retry_count += 1
-                    log_text += f"⚠️ 样本 {idx+1} API错误，重试 {retry_count}/{max_retries}: {str(e)}\n"
+                    log_text += f"⚠️ 样本 {idx+1} 处理错误，重试 {retry_count}/{max_retries}: {str(e)}\n"
             else:
                 results.append(f"[ERROR] Failed to process: {original_text}")
                 log_text += f"❌ 样本 {idx+1} 处理失败\n"
@@ -364,7 +450,7 @@ def clean_dataset(file_path, question_column, model_choice, temperature, max_sam
         lst_final = []
         for i in range(len(data_ori)):
             item = str(data_ori[i])
-            if '\n' in item:
+            if '\n' in item and model_choice != "WAC-GEC":
                 tmp_lines = [line.strip() for line in item.strip().split('\n') if line.strip()]
                 tmp_lines[-1] = lst_extracted[i]
                 lst_final.append('\n'.join(tmp_lines))
@@ -386,7 +472,8 @@ def clean_dataset(file_path, question_column, model_choice, temperature, max_sam
         
         original_filename = os.path.basename(file_path)
         base_name = original_filename.replace('.parquet', '')
-        output_filename = f"{base_name}-Denoising.parquet"
+        model_suffix = "WAC-GEC" if model_choice == "WAC-GEC" else "DeepSeek"
+        output_filename = f"{base_name}-Denoising-{model_suffix}.parquet"
         output_path = os.path.join(tempfile.gettempdir(), output_filename)
         
         df_cleaned.to_parquet(output_path, index=False)
@@ -394,6 +481,7 @@ def clean_dataset(file_path, question_column, model_choice, temperature, max_sam
         log_text += f"\n\n📊 处理完成！\n"
         log_text += f"{'='*50}\n"
         log_text += f"【基础统计】\n"
+        log_text += f"- 使用模型: {model_choice}\n"
         log_text += f"- 总样本数: {total}\n"
         log_text += f"- 成功处理: {total - error_count - unknown_count}\n"
         log_text += f"- 失败样本: {error_count}\n"
@@ -408,6 +496,10 @@ def clean_dataset(file_path, question_column, model_choice, temperature, max_sam
         log_text += f"📍 拼写错误密度（SED）:\n"
         log_text += f"   原始: {sed_original:.2f}% → 清洗后: {sed_cleaned:.2f}%\n"
         log_text += f"   ��化: {delta_sed:+.2f}% {'✅ 改善' if delta_sed < 0 else '⚠️ 增加'}\n"
+        
+        if model_choice == "WAC-GEC":
+            log_text += f"\n💡 注意: WAC-GEC仅修正空白符错误，不修正拼写和语法错误\n"
+        
         log_text += f"{'='*50}\n"
         
         # 生成带颜色的对比HTML
@@ -426,27 +518,39 @@ def clean_dataset(file_path, question_column, model_choice, temperature, max_sam
 ABOUT_TEXT = """
 ## 清洗流程说明
 
+### 支持的模型
+
+#### 1. DeepSeek-R1 (deepseek-r1-distill-llama-8b)
+- **功能**: 全面的语法、拼写、空格错误修正
+- **优势**: 综合性强，能处理多种类型的错误
+- **配置**: 需要在Space Settings中配置DEEPSEEK_API_KEY
+
+#### 2. WAC-GEC (Whitespace Correction)
+- **功能**: 专注于空白符错误纠正（多余空格、缺失空格等）
+- **优势**: 轻量级，无需API密钥，处理速度快
+- **限制**: 仅修正空白符错误，不处理拼写和语法问题
+- **适用场景**: 数据集中主要存在空白符异常的情况
+
 ### 核心算法
 
 1. **预处理 (process_sentence)**
    - 检测句子完整性
-   - 为不完整的句子添加标记 `___`
+   - 为不完整的句子添加标记 `___` (仅DeepSeek)
    - 保留多行文本格式
 
-2. **LLM清洗**
-   - 使用 DeepSeek API 进行语法、拼写、空格错误修正
-   - 重试机制：最多重试5次
-   - 稳定的 REST API 调用
+2. **模型清洗**
+   - **DeepSeek**: 使用API进行全面错误修正，重试机制最多5次
+   - **WAC-GEC**: 使用本地模型进行空白符纠正，重试机制最多3次
 
-3. **格式验证 (is_valid_output)**
+3. **格式验证**
    - 验证输出格式正确性
-   - 检查是否保留了 `___` 标记
+   - 检查标记保留情况
    - 长度合理性检查
 
 4. **后处理**
    - 提取清洗后的内容
    - 恢复原始多行格式
-   - 生成 `XXX-Denoising.parquet` 文件
+   - 生成带模型标识的Parquet文件
 
 ### 支持的数据集
 
@@ -466,9 +570,11 @@ ABOUT_TEXT = """
 ### 技术栈
 
 - **LLM**: DeepSeek API (deepseek-r1-distill-llama-8b)
+- **本地模型**: WAC-GEC (Whitespace Correction)
 - **前端**: Gradio 4.16.0
 - **数据处理**: Pandas + PyArrow (Parquet)
 - **差异对比**: Python difflib
+- **NLP工具**: spaCy, pyspellchecker
 - **API调用**: OpenAI SDK
 - **部署**: Hugging Face Spaces
 
@@ -477,9 +583,16 @@ ABOUT_TEXT = """
 - **WAR (Whitespace Anomaly Rate)**: 空白符异常率
 - **SED (Spelling Error Density)**: 拼写错误密度
 
+### 模型选择建议
+
+- **需要全面清洗**: 选择 DeepSeek-R1
+- **仅需修正空格**: 选择 WAC-GEC（更快，无需API）
+- **预算有限**: 优先使用 WAC-GEC
+- **追求最佳效果**: 使用 DeepSeek-R1
+
 ---
 
-**研究生毕业论文成果展示** | Powered by DeepSeek API
+**研究生毕业论文成果展示** | Powered by DeepSeek API & WAC-GEC
 """
 
 # ======================== Gradio界面 ========================
@@ -558,6 +671,11 @@ with demo:
         with gr.TabItem("🚀 BD-toolkit Demo", id=3):
             gr.Markdown("## BD-toolkit轻量化Demo展示")
             
+            # 模型可用性提示
+            model_status = "✅ DeepSeek-R1: " + ("已配置" if DEEPSEEK_API_KEY else "未配置API密钥")
+            model_status += " | ✅ WAC-GEC: " + ("可用" if WAC_GEC_AVAILABLE else "未安装")
+            gr.Markdown(f"**模型状态**: {model_status}")
+            
             with gr.Row():
                 with gr.Column():
                     file_input = gr.File(
@@ -574,7 +692,8 @@ with demo:
                     model_choice = gr.Dropdown(
                         choices=["deepseek-r1-distill-llama-8b", "WAC-GEC"],
                         value="deepseek-r1-distill-llama-8b",
-                        label="🤖 选择模型"
+                        label="🤖 选择模型",
+                        info="DeepSeek: 全面纠错 | WAC-GEC: 仅空白符纠正"
                     )
                     
                     temperature = gr.Slider(
@@ -582,7 +701,8 @@ with demo:
                         maximum=1.0,
                         value=0.1,
                         step=0.1,
-                        label="🌡️ Temperature"
+                        label="🌡️ Temperature",
+                        info="仅对DeepSeek生效"
                     )
                     
                     max_samples = gr.Slider(

requirements.txt

@@ -4,4 +4,5 @@ pandas
 pyarrow
 openai
 spacy
-pyspellchecker
+pyspellchecker
+whitespace-correction