Merge pull request #35 from SHUzhangshuo/main

feat(metrics): implement comprehensive memoization for TEDS algorithm

Author: e06084

tests/test_teds.py

@@ -4,8 +4,10 @@
 TEDS (树编辑距离相似性) 指标的综合测试
 """
 
+import re
 import unittest
 import sys
+import time
 from pathlib import Path
 
 # Add project root to path

webmainbench/metrics/teds_metrics.py

@@ -10,6 +10,7 @@
 
 from typing import Dict, Any, List, Optional
 import re
+from functools import lru_cache
 from lxml import etree, html
 from lxml.html import HtmlElement
 from bs4 import BeautifulSoup
@@ -87,9 +88,12 @@ def _calculate_score(self, predicted: Any, groundtruth: Any, **kwargs) -> Metric
                 "algorithm": "TEDS"
             }
 
+            # 清理缓存以释放内存
+            self._clear_memoization_cache()
+
             return MetricResult(
                 metric_name=self.name,
-                score=max(0.0, min(1.0, teds_score)),  # 删除多余的右括号
+                score=max(0.0, min(1.0, teds_score)),
                 details=details
             )
 
@@ -268,50 +272,64 @@ def _count_nodes(self, tree: Dict) -> int:
 
     def _tree_edit_distance(self, tree1: Dict, tree2: Dict) -> float:
         """
-        Calculate tree edit distance using dynamic programming.
+        Calculate tree edit distance using dynamic programming with memoization.
         
         This is a simplified version of the tree edit distance algorithm.
         For production use, consider using more sophisticated algorithms.
         """
+        # 初始化记忆化缓存
+        if not hasattr(self, '_memo_cache'):
+            self._memo_cache = {}
+        
+        # 创建缓存键
+        cache_key = (id(tree1), id(tree2))
+        if cache_key in self._memo_cache:
+            return self._memo_cache[cache_key]
+        
+        # 计算编辑距离
         if tree1 is None and tree2 is None:
-            return 0.0
-        if tree1 is None:
-            return float(self._count_nodes(tree2))
-        if tree2 is None:
-            return float(self._count_nodes(tree1))
-        
-        # Check if nodes are the same
-        if self._nodes_equal(tree1, tree2):
-            # Nodes are equal, calculate cost for children
-            children1 = tree1.get('children', [])
-            children2 = tree2.get('children', [])
-            
-            return self._list_edit_distance(children1, children2)
+            result = 0.0
+        elif tree1 is None:
+            result = float(self._count_nodes(tree2))
+        elif tree2 is None:
+            result = float(self._count_nodes(tree1))
         else:
-            # 检查结构是否相同（忽略文本内容）
-            if self._structure_equal(tree1, tree2):
-                # 结构相同，内容不同，使用内容编辑距离
-                content_distance = self._content_edit_distance(tree1, tree2)
-                children_cost = self._list_edit_distance(
-                    tree1.get('children', []), 
-                    tree2.get('children', [])
-                )
-                return content_distance + children_cost
-            else:
-                # 结构不同，使用原有的删除插入策略
-                # Option 1: Replace tree1 with tree2
-                cost_replace = 1.0 + self._list_edit_distance(
-                    tree1.get('children', []), 
-                    tree2.get('children', [])
-                )
+            # Check if nodes are the same
+            if self._nodes_equal(tree1, tree2):
+                # Nodes are equal, calculate cost for children
+                children1 = tree1.get('children', [])
+                children2 = tree2.get('children', [])
 
-                # Option 2: Delete tree1 and insert tree2
-                cost_delete_insert = (
-                    float(self._count_nodes(tree1)) + 
-                    float(self._count_nodes(tree2))
-                )
-                
-                return min(cost_replace, cost_delete_insert)
+                result = self._list_edit_distance(children1, children2)
+            else:
+                # 检查结构是否相同（忽略文本内容）
+                if self._structure_equal(tree1, tree2):
+                    # 结构相同，内容不同，使用内容编辑距离
+                    content_distance = self._content_edit_distance(tree1, tree2)
+                    children_cost = self._list_edit_distance(
+                        tree1.get('children', []), 
+                        tree2.get('children', [])
+                    )
+                    result = content_distance + children_cost
+                else:
+                    # 结构不同，使用原有的删除插入策略
+                    # Option 1: Replace tree1 with tree2
+                    cost_replace = 1.0 + self._list_edit_distance(
+                        tree1.get('children', []), 
+                        tree2.get('children', [])
+                    )
+                    
+                    # Option 2: Delete tree1 and insert tree2
+                    cost_delete_insert = (
+                        float(self._count_nodes(tree1)) + 
+                        float(self._count_nodes(tree2))
+                    )
+                    
+                    result = min(cost_replace, cost_delete_insert)
+        
+        # 缓存结果
+        self._memo_cache[cache_key] = result
+        return result
 
     def _structure_equal(self, tree1: Dict, tree2: Dict) -> bool:
         """Check if two trees have identical structure (same tag, attributes)"""
@@ -332,45 +350,77 @@ def _structure_equal(self, tree1: Dict, tree2: Dict) -> bool:
         return True
 
     def _content_edit_distance(self, tree1: Dict, tree2: Dict) -> float:
-        """Calculate content edit distance between two trees with same structure"""
-        if tree1['tag'] != tree2['tag']:
-            return 1.0  # 标签不同，惩罚1分
+        """Calculate content edit distance between two trees with same structure with memoization"""
+        # 初始化内容编辑距离缓存
+        if not hasattr(self, '_content_memo_cache'):
+            self._content_memo_cache = {}
 
-        # 如果是叶子节点（如td），计算文本内容的编辑距离
-        if tree1['tag'] == 'td' or not tree1.get('children'):
-            text1 = tree1.get('text', '')
-            text2 = tree2.get('text', '')
-            
-            if text1 == text2:
-                return 0.0  # 内容相同
-            
-            # 计算文本编辑距离
-            return self._text_edit_distance(text1, text2)
+        # 创建缓存键
+        cache_key = (id(tree1), id(tree2))
+        if cache_key in self._content_memo_cache:
+            return self._content_memo_cache[cache_key]
 
-        # 非叶子节点，递归计算子节点的内容编辑距离
-        children1 = tree1.get('children', [])
-        children2 = tree2.get('children', [])
+        if tree1['tag'] != tree2['tag']:
+            result = 1.0  # 标签不同，惩罚1分
+        else:
+            # 如果是叶子节点（如td），计算文本内容的编辑距离
+            if tree1['tag'] == 'td' or not tree1.get('children'):
+                text1 = tree1.get('text', '')
+                text2 = tree2.get('text', '')
+                
+                if text1 == text2:
+                    result = 0.0  # 内容相同
+                else:
+                    # 计算文本编辑距离
+                    result = self._text_edit_distance(text1, text2)
+            else:
+                # 非叶子节点，递归计算子节点的内容编辑距离
+                children1 = tree1.get('children', [])
+                children2 = tree2.get('children', [])
+                
+                result = self._list_content_edit_distance(children1, children2)
 
-        return self._list_content_edit_distance(children1, children2)
+        # 缓存结果
+        self._content_memo_cache[cache_key] = result
+        return result
 
     def _text_edit_distance(self, text1: str, text2: str) -> float:
-        """Calculate normalized edit distance between two text strings"""
-        if not text1 and not text2:
-            return 0.0
-        if not text1 or not text2:
-            return 1.0
+        """Calculate normalized edit distance between two text strings with memoization"""
+        # 初始化文本编辑距离缓存
+        if not hasattr(self, '_text_memo_cache'):
+            self._text_memo_cache = {}
+        
+        # 创建缓存键
+        cache_key = (text1, text2)
+        if cache_key in self._text_memo_cache:
+            return self._text_memo_cache[cache_key]
 
-        # 计算Levenshtein编辑距离
-        distance = self._levenshtein_distance(text1, text2)
-        max_len = max(len(text1), len(text2))
+        if not text1 and not text2:
+            result = 0.0
+        elif not text1 or not text2:
+            result = 1.0
+        else:
+            # 计算Levenshtein编辑距离
+            distance = self._levenshtein_distance(text1, text2)
+            max_len = max(len(text1), len(text2))
+            
+            # 返回归一化的编辑距离（0-1之间）
+            result = float(distance) / max_len if max_len > 0 else 0.0
 
-        # 返回归一化的编辑距离（0-1之间）
-        return float(distance) / max_len if max_len > 0 else 0.0
+        # 缓存结果
+        self._text_memo_cache[cache_key] = result
+        return result
 
     def _levenshtein_distance(self, s1: str, s2: str) -> int:
-        """Calculate Levenshtein distance between two strings"""
+        """Calculate Levenshtein distance between two strings with memoization"""
+        # 使用lru_cache装饰器进行缓存
+        return self._cached_levenshtein_distance(s1, s2)
+    
+    @lru_cache(maxsize=10000)
+    def _cached_levenshtein_distance(self, s1: str, s2: str) -> int:
+        """Cached version of Levenshtein distance calculation"""
         if len(s1) < len(s2):
-            return self._levenshtein_distance(s2, s1)
+            return self._cached_levenshtein_distance(s2, s1)
 
         if len(s2) == 0:
             return len(s1)
@@ -422,7 +472,16 @@ def _list_content_edit_distance(self, list1: List, list2: List) -> float:
         return dp[m][n]
 
     def _list_edit_distance(self, list1: List, list2: List) -> float:
-        """Calculate edit distance between two lists of trees (for structure comparison)"""
+        """Calculate edit distance between two lists of trees (for structure comparison) with memoization"""
+        # 初始化列表编辑距离缓存
+        if not hasattr(self, '_list_memo_cache'):
+            self._list_memo_cache = {}
+        
+        # 创建缓存键
+        cache_key = (id(list1), id(list2))
+        if cache_key in self._list_memo_cache:
+            return self._list_memo_cache[cache_key]
+        
         m, n = len(list1), len(list2)
 
         # Initialize DP matrix
@@ -453,7 +512,10 @@ def _list_edit_distance(self, list1: List, list2: List) -> float:
                     dp[i][j-1] + ins_cost       # Insert
                 )
 
-        return dp[m][n]
+        result = dp[m][n]
+        # 缓存结果
+        self._list_memo_cache[cache_key] = result
+        return result
 
     def _nodes_equal(self, node1: Dict, node2: Dict) -> bool:
         """Check if two tree nodes are equal."""
@@ -476,6 +538,19 @@ def _nodes_equal(self, node1: Dict, node2: Dict) -> bool:
                 return False
 
         return True
+    
+    def _clear_memoization_cache(self):
+        """清理所有记忆化缓存以释放内存"""
+        if hasattr(self, '_memo_cache'):
+            self._memo_cache.clear()
+        if hasattr(self, '_list_memo_cache'):
+            self._list_memo_cache.clear()
+        if hasattr(self, '_content_memo_cache'):
+            self._content_memo_cache.clear()
+        if hasattr(self, '_text_memo_cache'):
+            self._text_memo_cache.clear()
+        # 清理lru_cache
+        self._cached_levenshtein_distance.cache_clear()
 
 
 class StructureTEDSMetric(TEDSMetric):