diff --git a/sql_ml.py b/sql_ml.py index d0288bb..2a0fb19 100644 --- a/sql_ml.py +++ b/sql_ml.py @@ -5,6 +5,7 @@ from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.cluster import DBSCAN from sklearn.metrics.pairwise import cosine_similarity from tqdm import tqdm +import argparse # --- 规则和辅助函数 (与之前相同) --- PREDEFINED_RULES = [ @@ -16,6 +17,25 @@ PREDEFINED_RULES = [ # {'name': 'Sent GCash to Account', 'pattern': re.compile(r'^Sent GCash to .+? with account ending in \d+$')} ] +# 占位符到正则表达式的映射 +PLACEHOLDER_PATTERNS = { + '<金额>': r'([\d,.]+)', + '<付款人名称>': r'(.+?)', + '<收款人名称>': r'(.+?)', + '<付款人号码>': r'([\d\w\+\-\(\)]+)', + '<收款人号码>': r'([\d\w\+\-\(\)]+)', + '<银行4位数尾号>': r'(\d{4})', + '<参考号>': r'(.+?)', + '<交易单号>': r'(.+?)', + '<日期时间>': r'(.+?)', + '<日期>': r'(\d{2}-\d{2}-\d{4})', + '<时间>': r'(\d{1,2}:\d{2}\s[AP]M)', + '<消息>': r'(.+?)', + '<流水号>': r'(.+?)', + '<网络或发票号>': r'(.+?)', + '<交易类型>': r'(.+?)', +} + def normalize_text(text): # 模式 8: 从银行收款 (这条规则必须先运行) text = re.sub( @@ -135,6 +155,48 @@ def normalize_text(text): return text +def template_to_regex(template): + """ + 将模板转换为可用于提取参数的正则表达式 + """ + # 转义模板中的特殊字符,但保留占位符 + escaped_template = re.escape(template) + + # 将占位符映射到对应的正则表达式捕获组 + for placeholder, pattern in PLACEHOLDER_PATTERNS.items(): + escaped_placeholder = re.escape(placeholder) + # 替换占位符为对应的捕获组 + escaped_template = escaped_template.replace(escaped_placeholder, pattern) + + return escaped_template + +def extract_parameters(template, message): + """ + 从消息中提取参数值 + """ + # 生成正则表达式 + pattern = template_to_regex(template) + + # 匹配消息 + match = re.search(pattern, message) + + if match: + # 获取所有捕获组 + values = match.groups() + + # 获取模板中的占位符 + placeholders = re.findall(r'<[^>]+>', template) + + # 创建参数字典 + parameters = {} + for i, placeholder in enumerate(placeholders): + if i < len(values): + parameters[placeholder] = values[i] + + return parameters + + return {} + def run_dbscan_on_corpus(corpus, eps, min_samples, max_samples=10): if not corpus: return {} @@ -278,18 +340,90 @@ def extract_templates_iterative(input_file, output_file, rules, batch_size=1000, print(f"错误:找不到输入文件 '{input_file}'。") return +def extract_values_with_templates(input_file, template_file, output_file): + """ + 使用DBSCAN生成的模板从原始消息中提取参数值 + """ + print("--- 开始使用模板提取参数值 ---") + + # 读取模板 + templates = [] + with open(template_file, 'r', encoding='utf-8') as f: + for line in f: + template_data = json.loads(line) + templates.append(template_data['content']) + + print(f"已加载 {len(templates)} 个模板") + + # 从原始数据中提取值 + extracted_values = [] + + with open(input_file, 'r', encoding='utf-8') as f: + total_lines = sum(1 for _ in f) + + with open(input_file, 'r', encoding='utf-8') as f: + for line in tqdm(f, total=total_lines, desc="提取参数"): + try: + data = json.loads(line) + content = data.get('content', '') + + if not content: + continue + + # 尝试匹配每个模板 + for template in templates: + parameters = extract_parameters(template, content) + if parameters: + extracted_values.append({ + 'template': template, + 'message': content, + 'parameters': parameters + }) + # 找到匹配就跳出循环 + break + + except (json.JSONDecodeError, Exception): + continue + + # 保存提取的值 + with open(output_file, 'w', encoding='utf-8') as f: + for item in extracted_values: + json.dump(item, f, ensure_ascii=False) + f.write('\n') + + print(f"成功从 {len(extracted_values)} 条消息中提取参数,并保存到 '{output_file}'") + # --- 使用示例 --- # 假设您已经运行了上一个脚本,生成了 'content_filtered.jsonl' input_jsonl_file = 'content_filtered.jsonl' output_template_file = 'templates_iterative.txt' BATCH_PROCESSING_SIZE = 10000 # 可以根据你的内存和数据量调整 -extract_templates_iterative( - input_file=input_jsonl_file, - output_file=output_template_file, - rules=PREDEFINED_RULES, - batch_size=BATCH_PROCESSING_SIZE, - eps=0.4, - min_samples=5, # 稍微提高min_samples可以得到更可靠的模板 - max_samples_per_template=10 # 设置为正数以导出样本数据,0表示不导出 -) \ No newline at end of file +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Extract templates from GCash transaction data.') + parser.add_argument('--input_file', type=str, default=input_jsonl_file, help='Input JSONL file path') + parser.add_argument('--output_file', type=str, default=output_template_file, help='Output template file path') + parser.add_argument('--batch_size', type=int, default=BATCH_PROCESSING_SIZE, help='Batch processing size (data volume)') + parser.add_argument('--eps', type=float, default=0.4, help='DBSCAN eps parameter') + parser.add_argument('--min_samples', type=int, default=5, help='DBSCAN min_samples parameter') + parser.add_argument('--extract_values', action='store_true', help='Extract values using generated templates') + + args = parser.parse_args() + + if args.extract_values: + # 执行参数提取 + extract_values_with_templates( + input_file=args.input_file, + template_file=args.output_file, + output_file='extracted_parameters.jsonl' + ) + else: + # 执行模板提取 + extract_templates_iterative( + input_file=args.input_file, + output_file=args.output_file, + rules=PREDEFINED_RULES, + batch_size=args.batch_size, + eps=args.eps, + min_samples=args.min_samples + ) \ No newline at end of file