WAC-GEC

app.py

@@ -154,11 +154,22 @@ Corrected:"""
     # 生成修正
     inputs = gec_tokenizer(prompt, return_tensors="pt").to(gec_model.device)
     
+    # 检测设备类型以优化参数
+    is_cpu = str(gec_model.device) == "cpu" or not torch.cuda.is_available()
+    
+    # CPU优化参数：减少beam search和token长度
+    if is_cpu:
+        max_tokens = 256  # CPU模式减半
+        beams = 2         # 减少beam数量加速
+    else:
+        max_tokens = 512  # GPU模式保持
+        beams = 4
+    
     with torch.no_grad():
         outputs = gec_model.generate(
             **inputs,
-            max_new_tokens=512,
-            num_beams=4,
+            max_new_tokens=max_tokens,
+            num_beams=beams,
             do_sample=False,
             temperature=None,
             top_p=None
@@ -847,6 +858,11 @@ with demo:
             )
 
 if __name__ == "__main__":
+    # 可选：预加载模型（会增加启动时间）
+    # 如果想要预加载，取消下面两行的注释
+    print("🚀 预加载WAC-GEC模型...")
+    initialize_wac_gec()
+    
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,