# -*- coding: utf-8 -*- import os import logging import datetime from typing import Dict, Optional, Tuple from flask import Flask, request, jsonify from transformers import BertTokenizer, BertForSequenceClassification, AutoTokenizer, AutoModelForTokenClassification import torch from werkzeug.exceptions import BadRequest from ner_config import NERConfig, RepaymentNERConfig, IncomeNERConfig, FlightNERConfig, TrainNERConfig import re # 配置日志 logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[ logging.FileHandler('app.log'), logging.StreamHandler() ] ) logger = logging.getLogger(__name__) class ModelManager: def __init__(self): self.classifier_path = "./models/classifier" self.ner_path = "./models/ner_model/best_model" self.repayment_path = "./models/repayment_model/best_model" self.income_path = "./models/income_model/best_model" # self.flight_path = "./models/flight_model/best_model" # self.train_path = "./models/train_model/best_model" # 添加火车票模型路径 # 检查模型文件 self._check_model_files() # 加载模型 self.classifier_tokenizer, self.classifier_model = self._load_classifier() self.ner_tokenizer, self.ner_model = self._load_ner() self.repayment_tokenizer, self.repayment_model = self._load_repayment() self.income_tokenizer, self.income_model = self._load_income() # self.flight_tokenizer, self.flight_model = self._load_flight() # self.train_tokenizer, self.train_model = self._load_train() # 加载火车票模型 # 将模型设置为评估模式 self.classifier_model.eval() self.ner_model.eval() self.repayment_model.eval() self.income_model.eval() # self.flight_model.eval() # self.train_model.eval() # 设置火车票模型为评估模式 def _check_model_files(self): """检查模型文件是否存在""" if not os.path.exists(self.classifier_path): raise RuntimeError("分类模型文件不存在，请先运行训练脚本") if not os.path.exists(self.ner_path): raise RuntimeError("NER模型文件不存在，请先运行训练脚本") if not os.path.exists(self.repayment_path): raise RuntimeError("还款模型文件不存在，请先运行训练脚本") if not os.path.exists(self.income_path): raise RuntimeError("收入模型文件不存在，请先运行训练脚本") # if not os.path.exists(self.flight_path): # raise RuntimeError("航班模型文件不存在，请先运行训练脚本") # if not os.path.exists(self.train_path): # raise RuntimeError("火车票模型文件不存在，请先运行训练脚本") def _load_classifier(self) -> Tuple[BertTokenizer, BertForSequenceClassification]: """加载分类模型""" try: tokenizer = BertTokenizer.from_pretrained(self.classifier_path) model = BertForSequenceClassification.from_pretrained(self.classifier_path) return tokenizer, model except Exception as e: logger.error(f"加载分类模型失败: {str(e)}") raise def _load_ner(self) -> Tuple[AutoTokenizer, AutoModelForTokenClassification]: """加载NER模型""" try: tokenizer = AutoTokenizer.from_pretrained(self.ner_path) model = AutoModelForTokenClassification.from_pretrained(self.ner_path) return tokenizer, model except Exception as e: logger.error(f"加载NER模型失败: {str(e)}") raise def _load_repayment(self): """加载还款模型""" try: tokenizer = AutoTokenizer.from_pretrained(self.repayment_path) model = AutoModelForTokenClassification.from_pretrained(self.repayment_path) return tokenizer, model except Exception as e: logger.error(f"加载还款模型失败: {str(e)}") raise def _load_income(self): """加载收入模型""" try: tokenizer = AutoTokenizer.from_pretrained(self.income_path) model = AutoModelForTokenClassification.from_pretrained(self.income_path) return tokenizer, model except Exception as e: logger.error(f"加载收入模型失败: {str(e)}") raise def _load_flight(self): """加载航班模型""" try: tokenizer = AutoTokenizer.from_pretrained(self.flight_path) model = AutoModelForTokenClassification.from_pretrained(self.flight_path) return tokenizer, model except Exception as e: logger.error(f"加载航班模型失败: {str(e)}") raise def _load_train(self): """加载火车票模型""" try: tokenizer = AutoTokenizer.from_pretrained(self.train_path) model = AutoModelForTokenClassification.from_pretrained(self.train_path) return tokenizer, model except Exception as e: logger.error(f"加载火车票模型失败: {str(e)}") raise def classify_sms(self, text: str) -> Tuple[str, float]: """对短信进行分类，并返回置信度""" try: inputs = self.classifier_tokenizer( text, return_tensors="pt", truncation=True, max_length=64 ) with torch.no_grad(): outputs = self.classifier_model(**inputs) # 获取预测标签及其对应的概率 logits = outputs.logits probabilities = torch.softmax(logits, dim=1) pred_id = logits.argmax().item() confidence = probabilities[0, pred_id].item() # 获取预测标签的置信度 return self.classifier_model.config.id2label[pred_id], confidence except Exception as e: logger.error(f"短信分类失败: {str(e)}") raise def is_marketing_sms(self, text: str) -> bool: """判断是否为营销/广告类短信，采用评分系统""" # 特定字符串模式检查：直接匹配明显的营销/通知短信 marketing_patterns = [ # 百度类通知 r"百度智能云.*?尊敬的用户", r"百度.*?账户.*?tokens", r"AppBuilder.*?账户", r"账户有.*?免费额度", r".*?免费额度.*?过期", r"dwz\.cn\/[A-Za-z0-9]+" ] # 对特定模式直接判断 for pattern in marketing_patterns: if re.search(pattern, text): return True # 直接认为是营销短信 # 评分系统：根据短信内容特征进行评分，超过阈值判定为营销短信 score = 0 # 强营销特征关键词（高权重） strong_marketing_keywords = [ "有奖", "免费赠送", "抽奖", "中奖", "优惠券", "折扣券", "特价", "秒杀", "限时抢购", "促销", "推广", "广告", "代金券", "0元购", "tokens调用量" ] # 一般营销特征关键词（中等权重） general_marketing_keywords = [ "活动", "优惠", "折扣", "限时", "抢购", "特价", "promotion", "推广", "开业", "集点", "集赞", "关注", "公众号", "小程序", "注册有礼", "免费额度" ] # 弱营销特征关键词（低权重，可能出现在正常短信中） weak_marketing_keywords = [ "尊敬的用户", "尊敬的客户", "您好", "注册", "登录", "账户", "账号", "会员", "积分", "权益", "提醒", "即将", "有效期", "过期", "升级", "更新", "下载", "APP", "应用", "平台", "网址", "点击", "工单" ] # 短网址和链接（独立评估，结合其他特征判断） url_patterns = [ "dwz.cn", "t.cn", "短网址", "http://", "https://", "cmbt.cn" ] # 业务短信特征（用于反向识别，降低误判率） # 快递短信特征 express_keywords = [ "快递", "包裹", "取件码", "取件", "签收", "派送", "配送", "物流", "驿站", "在途", "揽收", "暂存", "已到达", "丰巢", "柜取件", "柜机" ] # 还款短信特征 repayment_keywords = [ "还款", "账单", "信用卡", "借款", "贷款", "逾期", "欠款", "最低还款", "应还金额", "到期还款", "还清", "应还", "还款日", "账单￥", "账单¥", "查账还款" ] # 收入短信特征 income_keywords = [ "收入", "转账", "入账", "到账", "支付", "工资", "报销", "余额", "成功收款", "收到", "款项" ] # 航班/火车票特征 travel_keywords = [ "航班", "航空", "飞机", "机票", "火车", "铁路", "列车", "车票", "出发", "抵达", "起飞", "登机", "候车", "检票" ] # 额外增加：通知类短信特征（通常不需要处理的短信） notification_keywords = [ "余额不足", "话费不足", "话费余额", "通讯费", "流量用尽", "流量不足", "停机", "恢复通话", "自动充值", "交费", "缴费", "消费提醒", "交易提醒", "动账", "短信通知", "验证码", "校验码", "安全码" ] # 运营商标识 telecom_keywords = [ "中国电信", "中国移动", "中国联通", "电信", "移动", "联通", "携号转网", "号码服务", "通讯服务", "189.cn", "10086", "10010" ] # 银行和金融机构标识 bank_keywords = [ "信用卡", "储蓄卡", "借记卡", "储蓄", "银联", "建设银行", "工商银行", "农业银行", "中国银行", "交通银行", "招商银行", "浦发银行", "民生银行", "兴业银行", "广发银行", "平安银行", "中信银行", "光大银行", "华夏银行", "邮储银行", "农商银行", "支付宝", "微信支付", "京东金融", "度小满", "陆金所" ] # 特殊情况检查：招商银行账单短信，不应被过滤 if ("招商银行" in text and ("账单" in text or "还款日" in text)) or "cmbt.cn" in text: if "还款" in text or "账单" in text or "消费卡" in text: return False # 是还款短信，不过滤 # 计算评分 # 首先检查业务短信特征，如果明确是业务短信，直接返回False has_express_feature = any(keyword in text for keyword in express_keywords) has_repayment_feature = any(keyword in text for keyword in repayment_keywords) has_income_feature = any(keyword in text for keyword in income_keywords) has_travel_feature = any(keyword in text for keyword in travel_keywords) # 检查是否为百度通知 is_baidu_notification = "百度" in text and "尊敬的用户" in text if is_baidu_notification: return True # 百度通知应被过滤 # 如果短信中包含多个业务关键词（≥2个），很可能是重要的业务短信 business_score = (has_express_feature + has_repayment_feature + has_income_feature + has_travel_feature) if business_score >= 2 and not is_baidu_notification: return False # 多个业务特征同时存在，不太可能是营销短信 # 检查强营销特征 for keyword in strong_marketing_keywords: if keyword in text: score += 3 # 检查一般营销特征 for keyword in general_marketing_keywords: if keyword in text: score += 2 # 检查弱营销特征 for keyword in weak_marketing_keywords: if keyword in text: score += 1 # 检查URL特征（结合是否存在业务特征） has_url = any(pattern in text for pattern in url_patterns) # 降低业务特征短信的营销判定分数 if has_express_feature and not is_baidu_notification: score -= 3 # 快递特征明显减分 if has_repayment_feature: score -= 3 # 还款特征明显减分 if has_income_feature: score -= 2 # 收入特征减分 if has_travel_feature: score -= 2 # 旅行特征减分 # 检查通知类短信特征（但不包括重要的业务短信） if not has_express_feature and not has_repayment_feature: # 确保不是快递和还款短信 notification_count = sum(1 for keyword in notification_keywords if keyword in text) if notification_count >= 2: # 需要至少2个通知关键词才判定 score += notification_count # 增加判定为营销/通知短信的可能性 # 检查运营商和银行标识（结合其他特征判断） has_telecom_feature = any(keyword in text for keyword in telecom_keywords) has_bank_feature = any(keyword in text for keyword in bank_keywords) # URL的评分处理 if has_url: if (has_express_feature or has_repayment_feature or has_income_feature or has_travel_feature) and not is_baidu_notification: # URL在业务短信中可能是正常的追踪链接，不增加评分 pass else: # 纯URL且无业务特征，可能是营销短信 score += 2 # 特殊情况：运营商余额通知 if has_telecom_feature and "余额" in text and not has_income_feature: score += 2 # 设置判定阈值 threshold = 4 # 需要至少4分才判定为营销短信 return score >= threshold def is_notification_sms(self, text: str) -> bool: """判断是否为通知类短信（如银行交易通知、运营商提醒等）""" # 银行交易通知特征（不包括还款提醒） bank_transaction_patterns = [ r"您尾号\d+的.+消费", r"您.+账户消费[\d,.]+元", r"交易[\d,.]+元", r"支付宝.+消费", r"微信支付.+消费", r"\d{1,2}月\d{1,2}日\d{1,2}[:：]\d{1,2}消费", r"银行卡([支付|消费|扣款])" ] # 排除规则：包含以下关键词的短信不应被判定为通知短信 business_keywords = [ # 还款关键词 "还款", "账单", "应还", "到期还款", "还款日", "最低还款", "账单￥", "账单¥", "查账还款", # 快递关键词 "快递", "包裹", "取件码", "取件", "签收", "派送", "配送", # 收入关键词 "收入", "转账", "入账", "到账", "支付成功", "工资" ] # 运营商余额通知特征 telecom_balance_patterns = [ r"余额[不足|低于][\d,.]+元", r"话费[不足|仅剩][\d,.]+元", r"流量[不足|即将用尽]", r"[电信|移动|联通].+余额", r"[停机|停号]提醒", r"为了保障您的正常通讯", ] # 首先检查是否包含业务关键词，有则不应判定为通知短信 for keyword in business_keywords: if keyword in text: return False # 包含业务关键词，不是需要过滤的通知短信 # 检查银行交易通知模式 for pattern in bank_transaction_patterns: if re.search(pattern, text): logger.debug(f"识别到银行交易通知短信：{text[:30]}...") return True # 检查运营商余额通知模式 for pattern in telecom_balance_patterns: if re.search(pattern, text): logger.debug(f"识别到运营商余额通知短信：{text[:30]}...") return True return False def extract_entities(self, text: str) -> Dict[str, Optional[str]]: """提取文本中的实体""" try: # 初始化结果字典 result = { "post": None, # 快递公司 "company": None, # 寄件公司 "address": None, # 地址 "pickup_code": None, # 取件码 "time": None # 添加时间字段 } # 第一阶段：直接从文本中提取取件码 pickup_code = self.extract_pickup_code_from_text(text) if pickup_code: result["pickup_code"] = pickup_code # 第二阶段：使用NER模型提取其他实体 inputs = self.ner_tokenizer( text, return_tensors="pt", truncation=True, max_length=64 ) with torch.no_grad(): outputs = self.ner_model(**inputs) predictions = torch.argmax(outputs.logits, dim=2) tokens = self.ner_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]) tags = [self.ner_model.config.id2label[p] for p in predictions[0].numpy()] # 解析实体 current_entity = None entity_start = None for i, (token, tag) in enumerate(zip(tokens, tags)): # 跳过取件码实体，因为我们已经单独处理了 if tag.startswith("B-") and tag[2:] != "PICKUP_CODE": # 保存之前的实体 if current_entity and current_entity["type"] != "PICKUP_CODE": key = current_entity["type"].lower() result[key] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() # 开始新实体 current_entity = {"type": tag[2:], "text": token} entity_start = i elif tag.startswith("I-") and current_entity and tag[2:] == current_entity["type"] and current_entity["type"] != "PICKUP_CODE": current_entity["text"] += token elif tag == "O" or tag.startswith("B-"): # 结束当前实体 if current_entity and current_entity["type"] != "PICKUP_CODE": key = current_entity["type"].lower() result[key] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() current_entity = None # 处理最后一个实体 if current_entity and current_entity["type"] != "PICKUP_CODE": key = current_entity["type"].lower() result[key] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() # 如果第一阶段没有提取到取件码，使用NER模型的结果 if not result["pickup_code"]: # 重新解析一次，只提取取件码 current_entity = None for i, (token, tag) in enumerate(zip(tokens, tags)): if tag.startswith("B-") and tag[2:] == "PICKUP_CODE": current_entity = {"type": tag[2:], "text": token} elif tag.startswith("I-") and current_entity and tag[2:] == "PICKUP_CODE": current_entity["text"] += token elif (tag == "O" or tag.startswith("B-")) and current_entity: result["pickup_code"] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() current_entity = None # 处理最后一个实体 if current_entity and current_entity["type"] == "PICKUP_CODE": result["pickup_code"] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() # 后处理：清理和验证取件码 if result["pickup_code"]: # 清理取件码 code = result["pickup_code"] # 移除取件码后的额外文字 for word in ["的", "来", "到", "取", "件", "码"]: if word in code: code = code[:code.index(word)] # 只保留字母、数字和连字符 code = ''.join(c for c in code if c.isalnum() or c == "-") # 确保格式正确 parts = code.split("-") valid_parts = [] for part in parts: if part and any(c.isalnum() for c in part): valid_parts.append(part) if valid_parts: result["pickup_code"] = "-".join(valid_parts) else: result["pickup_code"] = None # 清理公司名称 if result["company"]: company = result["company"] invalid_words = ["包裹", "快递", "到"] for word in invalid_words: if company.endswith(word): company = company[:-len(word)] result["company"] = company.strip() # 清理地址 if result["address"]: address = result["address"] invalid_suffixes = ["，请尽快取件", "，询", "请尽快取件"] for suffix in invalid_suffixes: if address.endswith(suffix): address = address[:-len(suffix)] result["address"] = address.strip() return result except Exception as e: logger.error(f"实体提取失败: {str(e)}") raise def extract_pickup_code_from_text(self, text: str) -> Optional[str]: """直接从文本提取取件码""" # 先尝试直接匹配完整的取件码格式 pickup_patterns = [ # 带有上下文的取件码模式（包括字母数字组合） r'凭\s*([A-Za-z0-9]{1,3}-\d{1,2}-\d{1,6})\s*(来|到)?', # 匹配"凭xx-xx-xxxx"格式 r'码\s*([A-Za-z0-9]{1,3}-\d{1,2}-\d{1,6})', # 匹配"码xx-xx-xxxx"格式 r'提货码\s*([A-Za-z0-9]{1,3}-\d{1,2}-\d{1,6})', # 匹配"提货码xx-xx-xxxx"格式 r'取件码\s*([A-Za-z0-9]{1,3}-\d{1,2}-\d{1,6})', # 匹配"取件码xx-xx-xxxx"格式 # 分组提取模式，对有上下文的取件码 r'凭\s*([A-Za-z0-9]{1,3})-(\d{1,2})-(\d{1,6})\s*(来|到)?', # 匹配"凭xx-xx-xxxx" r'码\s*([A-Za-z0-9]{1,3})-(\d{1,2})-(\d{1,6})', # 匹配"码xx-xx-xxxx" r'提货码\s*([A-Za-z0-9]{1,3})-(\d{1,2})-(\d{1,6})', # 匹配"提货码xx-xx-xxxx" r'取件码\s*([A-Za-z0-9]{1,3})-(\d{1,2})-(\d{1,6})', # 匹配"取件码xx-xx-xxxx" # 独立的取件码模式 r'([A-Za-z0-9]{1,3}-\d{1,2}-\d{1,6})', # 匹配独立的xx-xx-xxxx格式，支持字母 r'(\d{1,2}-\d{1,2}-\d{4,6})', # 匹配独立的xx-xx-xxxx格式，纯数字 ] # 首先尝试整体匹配 for pattern in pickup_patterns[:4]: # 前4个是整体匹配模式 matches = re.findall(pattern, text) if matches: return matches[0] if isinstance(matches[0], str) else matches[0][0] # 然后尝试分组匹配 for pattern in pickup_patterns[4:8]: # 中间4个是分组匹配模式 matches = re.findall(pattern, text) if matches: # 分组匹配，第一部分可以是字母数字组合 match = matches[0] if len(match) >= 3: # 确保第2、3部分是数字 if match[1].isdigit() and match[2].isdigit(): return f"{match[0]}-{match[1]}-{match[2]}" # 最后尝试独立模式 for pattern in pickup_patterns[8:10]: # 最后2个是独立匹配模式 matches = re.findall(pattern, text) if matches: return matches[0] # 特殊处理包含"取件码"的文本 code_indicators = ["取件码", "提货码"] for indicator in code_indicators: if indicator in text: idx = text.find(indicator) # 检查紧跟着的文本 search_text = text[idx:idx+35] # 扩大搜索范围 # 尝试匹配"取件码A8-1-15"这样的格式 special_match = re.search(r'[码]\s*([A-Za-z0-9]+[-\s]+\d+[-\s]+\d+)', search_text) if special_match: # 规范化格式 code = special_match.group(1) # 将空格替换为连字符，确保格式一致 code = re.sub(r'\s+', '-', code) # 确保连字符格式正确 code = re.sub(r'-+', '-', code) return code # 特殊处理"凭"后面的数字序列 if "凭" in text: pickup_index = text.find("凭") if pickup_index != -1: # 在"凭"之后的25个字符内寻找取件码 search_text = text[pickup_index:pickup_index+35] # 尝试提取形如"凭A8-1-15"或"凭22-4-1111"的格式 alpha_num_match = re.search(r'凭\s*([A-Za-z0-9]+)[-\s]+(\d+)[-\s]+(\d+)', search_text) if alpha_num_match: return f"{alpha_num_match.group(1)}-{alpha_num_match.group(2)}-{alpha_num_match.group(3)}" # 提取所有字母数字序列 parts = re.findall(r'[A-Za-z0-9]+', search_text) if len(parts) >= 3: # 组合前三个部分 return f"{parts[0]}-{parts[1]}-{parts[2]}" # 查找形如"A8-1-15"的取件码 alpha_num_codes = re.findall(r'([A-Za-z]\d+)-(\d+)-(\d+)', text) if alpha_num_codes: match = alpha_num_codes[0] return f"{match[0]}-{match[1]}-{match[2]}" return None def extract_repayment_entities(self, text: str) -> Dict[str, Optional[str]]: """提取还款相关实体""" try: result = { "bank": None, # 还款主体 "type": None, # 还款类型 "amount": None, # 还款金额 "date": None, # 还款日期 "number": None, # 账号尾号 "min_amount": None # 最低还款金额 } inputs = self.repayment_tokenizer( text, return_tensors="pt", truncation=True, max_length=RepaymentNERConfig.MAX_LENGTH ) with torch.no_grad(): outputs = self.repayment_model(**inputs) predictions = torch.argmax(outputs.logits, dim=2) tokens = self.repayment_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]) tags = [self.repayment_model.config.id2label[p] for p in predictions[0].numpy()] def clean_amount(amount_text: str, context: str) -> Optional[str]: """清理和标准化金额""" if not amount_text: return None # 尝试直接在上下文中使用正则表达式查找更完整的金额 # 如果在同一句话里有类似"应还金额5,800元"这样的模式 amount_match = re.search(r'(?:应还|还款)?金额([\d,]+\.?\d*)(?:元|块钱|块|万元|万)?', context) if amount_match: return amount_match.group(1) # 直接返回匹配到的金额，保留原始格式 # 尝试查找最低还款金额 min_amount_match = re.search(r'最低还款([\d,]+\.?\d*)(?:元|块钱|块|万元|万)?', context) if min_amount_match and "MIN_CODE" in current_entity["type"]: return min_amount_match.group(1) # 直接返回匹配到的最低还款金额，保留原始格式 # 在上下文中查找完整金额 amount_index = context.find(amount_text) if amount_index != -1: # 扩大搜索范围，查找完整金额 search_start = max(0, amount_index - 10) search_end = min(len(context), amount_index + len(amount_text) + 10) search_text = context[search_start:search_end] # 使用正则表达式查找金额 amount_pattern = re.compile(r'([\d,]+\.?\d*)(?:元|块钱|块|万元|万)?') matches = list(amount_pattern.finditer(search_text)) if matches: # 选择最接近的匹配结果 best_match = None min_distance = float('inf') for match in matches: distance = abs(match.start() - (amount_index - search_start)) if distance < min_distance: min_distance = distance best_match = match.group(1) # 只取数字部分，保留逗号 if best_match: return best_match # 如果上述方法都没找到，则保留原始提取结果但验证其有效性 # 移除货币符号和无效词 for symbol in RepaymentNERConfig.AMOUNT_CONFIG['currency_symbols']: amount_text = amount_text.replace(symbol, '') for word in RepaymentNERConfig.AMOUNT_CONFIG['invalid_words']: amount_text = amount_text.replace(word, '') # 验证金额有效性 clean_amount = amount_text.replace(',', '') try: value = float(clean_amount) if value > 0: # 返回原始格式 return amount_text except ValueError: pass return None # 实体提取 current_entity = None entities = { "BANK": [], "TYPE": [], "PICKUP_CODE": [], "DATE": [], "NUMBER": [], "MIN_CODE": [] } for token, tag in zip(tokens, tags): if tag.startswith("B-"): if current_entity: entities[current_entity["type"]].append(current_entity["text"]) current_entity = {"type": tag[2:], "text": token.replace("##", "")} elif tag.startswith("I-") and current_entity and tag[2:] == current_entity["type"]: current_entity["text"] += token.replace("##", "") else: if current_entity: entities[current_entity["type"]].append(current_entity["text"]) current_entity = None if current_entity: entities[current_entity["type"]].append(current_entity["text"]) # 处理银行名称 if entities["BANK"]: bank_parts = [] seen = set() for bank in entities["BANK"]: bank = bank.strip() if bank and bank not in seen: bank_parts.append(bank) seen.add(bank) bank = "".join(bank_parts) if len(bank) <= RepaymentNERConfig.MAX_ENTITY_LENGTH["BANK"]: result["bank"] = bank # 处理还款类型 if entities["TYPE"]: type_parts = [] seen = set() for type_ in entities["TYPE"]: type_ = type_.strip() if type_ and type_ not in seen: type_parts.append(type_) seen.add(type_) type_ = "".join(type_parts) if len(type_) <= RepaymentNERConfig.MAX_ENTITY_LENGTH["TYPE"]: result["type"] = type_ # 处理账号尾号 if entities["NUMBER"]: number = "".join(c for c in "".join(entities["NUMBER"]) if c.isdigit()) if number and len(number) <= RepaymentNERConfig.MAX_ENTITY_LENGTH["NUMBER"]: result["number"] = number # 处理日期 if entities["DATE"]: date = "".join(entities["DATE"]) date = ''.join(c for c in date if c.isdigit() or c in ['年', '月', '日', '-']) if date: result["date"] = date # 处理金额 # 尝试匹配带￥符号的账单金额模式 amount_match = re.search(r'账单￥([\d,]+\.?\d*)', text) if not amount_match: # 尝试匹配带¥符号的账单金额模式 amount_match = re.search(r'账单¥([\d,]+\.?\d*)', text) if not amount_match: # 尝试匹配一般金额模式 amount_match = re.search(r'(?:应还|还款)?金额([\d,]+\.?\d*)(?:元|块钱|块|万元|万)?', text) if amount_match: amount = amount_match.group(1) # 保留原始格式(带逗号) # 验证金额有效性 try: value = float(amount.replace(',', '')) if value > 0: result["amount"] = amount except ValueError: pass # 如果正则没有匹配到，使用NER结果 if not result["amount"]: amount_candidates = [] # 从识别的实体中获取 for amount in entities["PICKUP_CODE"]: cleaned_amount = clean_amount(amount, text) if cleaned_amount: try: value = float(cleaned_amount.replace(',', '')) amount_candidates.append((cleaned_amount, value)) except ValueError: continue # 如果还是没有找到，尝试从文本中提取 if not amount_candidates: # 使用多个正则表达式匹配不同格式的金额 # 1. 匹配带￥符号格式 matches = list(re.finditer(r'￥([\d,]+\.?\d*)', text)) # 2. 匹配带¥符号格式 matches.extend(list(re.finditer(r'¥([\d,]+\.?\d*)', text))) # 3. 匹配一般金额格式 matches.extend(list(re.finditer(r'([\d,]+\.?\d*)(?:元|块钱|块|万元|万)', text))) for match in matches: amount_text = match.group(1) # 获取数字部分，保留逗号 try: value = float(amount_text.replace(',', '')) amount_candidates.append((amount_text, value)) except ValueError: continue # 选择最大的有效金额 if amount_candidates: result["amount"] = max(amount_candidates, key=lambda x: x[1])[0] # 处理最低还款金额 # 先尝试使用正则表达式直接匹配最低还款金额 min_amount_match = re.search(r'最低还款([\d,]+\.?\d*)(?:元|块钱|块|万元|万)?', text) if min_amount_match: min_amount = min_amount_match.group(1) # 保留原始格式(带逗号) # 验证金额有效性 try: value = float(min_amount.replace(',', '')) if value > 0: result["min_amount"] = min_amount except ValueError: pass # 如果正则没有匹配到，使用NER结果 if not result["min_amount"] and entities["MIN_CODE"]: for amount in entities["MIN_CODE"]: cleaned_amount = clean_amount(amount, text) if cleaned_amount: result["min_amount"] = cleaned_amount break return result except Exception as e: logger.error(f"还款实体提取失败: {str(e)}") raise def extract_income_entities(self, text: str) -> Dict[str, Optional[str]]: """提取收入相关实体""" try: result = { "bank": None, # 银行名称 "number": None, # 账号尾号 "datetime": None, # 交易时间 "amount": None, # 收入金额 "balance": None # 余额 } inputs = self.income_tokenizer( text, return_tensors="pt", truncation=True, max_length=IncomeNERConfig.MAX_LENGTH ) with torch.no_grad(): outputs = self.income_model(**inputs) predictions = torch.argmax(outputs.logits, dim=2) tokens = self.income_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]) tags = [self.income_model.config.id2label[p] for p in predictions[0].numpy()] def clean_amount(amount_text: str, context: str) -> Optional[str]: """清理和标准化金额""" if not amount_text: return None # 在上下文中查找完整金额 amount_index = context.find(amount_text) if amount_index != -1: # 扩大搜索范围，查找完整金额 search_start = max(0, amount_index - 10) search_end = min(len(context), amount_index + len(amount_text) + 10) search_text = context[search_start:search_end] # 使用更精确的正则表达式查找金额模式，支持带逗号的金额 amount_pattern = re.compile(r'(\d{1,3}(?:,\d{3})*(?:\.\d{1,2})?|\d+(?:\.\d{1,2})?)') matches = list(amount_pattern.finditer(search_text)) # 找到最接近且最长的完整金额 best_match = None min_distance = float('inf') max_length = 0 target_pos = amount_index - search_start for match in matches: match_pos = match.start() distance = abs(match_pos - target_pos) match_text = match.group(1) if len(match_text) > max_length or (len(match_text) == max_length and distance < min_distance): try: value = float(match_text) if value > 0: best_match = match_text min_distance = distance max_length = len(match_text) except ValueError: continue if best_match: amount_text = best_match # 移除货币符号和无效词 for symbol in IncomeNERConfig.AMOUNT_CONFIG['currency_symbols']: amount_text = amount_text.replace(symbol, '') for word in IncomeNERConfig.AMOUNT_CONFIG['invalid_words']: amount_text = amount_text.replace(word, '') # 处理金额中的逗号 amount_text = amount_text.replace(',', '') try: value = float(amount_text) if '.' in amount_text: decimal_places = len(amount_text.split('.')[1]) return f"{value:.{decimal_places}f}" return str(int(value)) except ValueError: pass return None # 实体提取 current_entity = None entities = { "BANK": [], "NUMBER": [], "DATATIME": [], "PICKUP_CODE": [], "BALANCE": [] } for token, tag in zip(tokens, tags): if tag.startswith("B-"): if current_entity: entities[current_entity["type"]].append(current_entity["text"]) current_entity = {"type": tag[2:], "text": token.replace("##", "")} elif tag.startswith("I-") and current_entity and tag[2:] == current_entity["type"]: current_entity["text"] += token.replace("##", "") else: if current_entity: entities[current_entity["type"]].append(current_entity["text"]) current_entity = None if current_entity: entities[current_entity["type"]].append(current_entity["text"]) # 处理银行名称 if entities["BANK"]: bank_parts = [] seen = set() for bank in entities["BANK"]: bank = bank.strip() if bank and bank not in seen: bank_parts.append(bank) seen.add(bank) bank = "".join(bank_parts) if len(bank) <= IncomeNERConfig.MAX_ENTITY_LENGTH["BANK"]: result["bank"] = bank # 处理账号尾号 if entities["NUMBER"]: number = "".join(c for c in "".join(entities["NUMBER"]) if c.isdigit()) if number and len(number) <= IncomeNERConfig.MAX_ENTITY_LENGTH["NUMBER"]: result["number"] = number # 处理交易时间 if entities["DATATIME"]: datetime = "".join(entities["DATATIME"]) datetime = ''.join(c for c in datetime if c.isdigit() or c in ['年', '月', '日', '时', '分', ':', '-']) if datetime: result["datetime"] = datetime # 处理收入金额 # 先尝试使用正则表达式直接匹配收入金额，包括"收入金额"格式 amount_match = re.search(r'收入金额([\d,]+\.?\d*)元', text) if not amount_match: # 尝试匹配一般收入格式 amount_match = re.search(r'收入([\d,]+\.?\d*)元', text) if amount_match: amount = amount_match.group(1) # 保留原始格式(带逗号) # 验证金额有效性 try: value = float(amount.replace(',', '')) if value > 0: result["amount"] = amount except ValueError: pass # 如果正则没有匹配到，继续尝试NER结果 if not result["amount"]: amount_candidates = [] # 首先从识别的实体中获取 for amount in entities["PICKUP_CODE"]: cleaned_amount = clean_amount(amount, text) if cleaned_amount: try: value = float(cleaned_amount) amount_candidates.append((cleaned_amount, value)) except ValueError: continue # 如果没有找到有效金额，直接从文本中尝试提取 if not amount_candidates: # 尝试多种模式匹配金额 # 1. 匹配"收入金额xxx元"模式 matches = list(re.finditer(r'收入金额([\d,]+\.?\d*)元', text)) # 2. 匹配"收入xxx元"模式 matches.extend(list(re.finditer(r'收入([\d,]+\.?\d*)元', text))) # 3. 匹配带元结尾的金额 matches.extend(list(re.finditer(r'([0-9,]+\.[0-9]+)元', text))) # 4. 匹配普通数字(可能是余额)，但排除已识别为余额的金额 if "余额" in text: balance_match = re.search(r'余额([\d,]+\.?\d*)元', text) if balance_match: balance_value = balance_match.group(1) # 只匹配不等于余额的金额 all_numbers = re.finditer(r'(\d{1,3}(?:,\d{3})*(?:\.\d{1,2})?|\d+(?:\.\d{1,2})?)', text) for match in all_numbers: if match.group(1) != balance_value: matches.append(match) else: matches.extend(list(re.finditer(r'(\d{1,3}(?:,\d{3})*(?:\.\d{1,2})?|\d+(?:\.\d{1,2})?)', text))) for match in matches: amount_text = match.group(1) try: value = float(amount_text.replace(',', '')) amount_candidates.append((amount_text, value)) except ValueError: continue # 从金额候选中排除已识别的余额值 if result["balance"]: try: balance_value = float(result["balance"].replace(',', '')) amount_candidates = [(text, value) for text, value in amount_candidates if abs(value - balance_value) > 0.01] except ValueError: pass # 选择适当的金额作为收入 if amount_candidates: has_income_amount_keyword = "收入金额" in text if has_income_amount_keyword: # 查找"收入金额"附近的数字 idx = text.find("收入金额") if idx != -1: closest_amount = None min_distance = float('inf') for amount_text, value in amount_candidates: # 找到这个数字在原文中的位置 amount_idx = text.find(amount_text) if amount_idx != -1: distance = abs(amount_idx - idx) if distance < min_distance: min_distance = distance closest_amount = amount_text if closest_amount: result["amount"] = closest_amount else: # 如果无法找到最近的金额，使用最大金额策略 result["amount"] = max(amount_candidates, key=lambda x: x[1])[0] else: # 如果没有"收入金额"关键词，则使用最大金额策略 result["amount"] = max(amount_candidates, key=lambda x: x[1])[0] # 处理余额 # 先尝试使用正则表达式直接匹配余额 balance_match = re.search(r'余额([\d,]+\.?\d*)元', text) if balance_match: balance = balance_match.group(1) # 保留原始格式(带逗号) # 验证金额有效性 try: value = float(balance.replace(',', '')) if value > 0: result["balance"] = balance except ValueError: pass # 如果正则没有匹配到，使用NER结果 if not result["balance"] and entities["BALANCE"]: for amount in entities["BALANCE"]: cleaned_amount = clean_amount(amount, text) if cleaned_amount: result["balance"] = cleaned_amount break return result except Exception as e: logger.error(f"收入实体提取失败: {str(e)}") raise def extract_flight_entities(self, text: str) -> Dict[str, Optional[str]]: """提取航班相关实体""" try: # 初始化结果字典 result = { "flight": None, # 航班号 "company": None, # 航空公司 "start": None, # 出发地 "end": None, # 目的地 "date": None, # 日期 "time": None, # 时间 "departure_time": None, # 起飞时间 "arrival_time": None, # 到达时间 "ticket_num": None, # 机票号码 "seat": None # 座位等信息 } # 使用NER模型提取实体 inputs = self.flight_tokenizer( text, return_tensors="pt", truncation=True, max_length=FlightNERConfig.MAX_LENGTH ) with torch.no_grad(): outputs = self.flight_model(**inputs) predictions = torch.argmax(outputs.logits, dim=2) tokens = self.flight_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]) tags = [self.flight_model.config.id2label[p] for p in predictions[0].numpy()] # 解析实体 current_entity = None for token, tag in zip(tokens, tags): if tag.startswith("B-"): if current_entity: entity_type = current_entity["type"].lower() result[entity_type] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() current_entity = {"type": tag[2:], "text": token} elif tag.startswith("I-") and current_entity and tag[2:] == current_entity["type"]: current_entity["text"] += token else: if current_entity: entity_type = current_entity["type"].lower() result[entity_type] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() current_entity = None # 处理最后一个实体 if current_entity: entity_type = current_entity["type"].lower() result[entity_type] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() # 处理航班号格式 if result["flight"]: flight_no = result["flight"].upper() # 清理航班号，只保留字母和数字 flight_no = ''.join(c for c in flight_no if c.isalnum()) # 验证航班号格式 valid_pattern = re.compile(FlightNERConfig.FLIGHT_CONFIG['pattern']) if valid_pattern.match(flight_no): result["flight"] = flight_no else: # 尝试修复常见错误 if len(flight_no) >= FlightNERConfig.FLIGHT_CONFIG['min_length'] and flight_no[:2].isalpha() and flight_no[2:].isdigit(): result["flight"] = flight_no else: result["flight"] = None # 清理日期格式 if result["date"]: date_str = result["date"] # 保留数字和常见日期分隔符 date_str = ''.join(c for c in date_str if c.isdigit() or c in ['年', '月', '日', '-', '/', '.']) result["date"] = date_str # 清理时间格式 for time_field in ["time", "departure_time", "arrival_time"]: if result[time_field]: time_str = result[time_field] # 保留数字和常见时间分隔符 time_str = ''.join(c for c in time_str if c.isdigit() or c in [':', '时', '分', '点']) result[time_field] = time_str # 处理机票号码 if result["ticket_num"]: ticket_num = result["ticket_num"] # 清理机票号码，只保留字母和数字 ticket_num = ''.join(c for c in ticket_num if c.isalnum()) result["ticket_num"] = ticket_num # 处理座位信息 if result["seat"]: seat_str = result["seat"] # 移除可能的额外空格和特殊字符 seat_str = seat_str.replace(" ", "").strip() result["seat"] = seat_str return result except Exception as e: logger.error(f"航班实体提取失败: {str(e)}") raise def extract_train_entities(self, text: str) -> Dict[str, Optional[str]]: """提取火车票相关实体""" try: # 初始化结果字典 result = { "company": None, # 12306 "trips": None, # 车次 "start": None, # 出发站 "end": None, # 到达站 "date": None, # 日期 "time": None, # 时间 "seat": None, # 座位等信息 "name": None # 用户姓名 } # 使用NER模型提取实体 inputs = self.train_tokenizer( text, return_tensors="pt", truncation=True, max_length=TrainNERConfig.MAX_LENGTH ) with torch.no_grad(): outputs = self.train_model(**inputs) predictions = torch.argmax(outputs.logits, dim=2) tokens = self.train_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]) tags = [self.train_model.config.id2label[p] for p in predictions[0].numpy()] # 解析实体 current_entity = None for token, tag in zip(tokens, tags): if tag.startswith("B-"): if current_entity: entity_type = current_entity["type"].lower() result[entity_type] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() current_entity = {"type": tag[2:], "text": token} elif tag.startswith("I-") and current_entity and tag[2:] == current_entity["type"]: current_entity["text"] += token else: if current_entity: entity_type = current_entity["type"].lower() result[entity_type] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() current_entity = None # 处理最后一个实体 if current_entity: entity_type = current_entity["type"].lower() result[entity_type] = current_entity["text"].replace("[UNK]", "").replace("##", "").strip() # 处理公司名称，通常为12306 if result["company"]: company = result["company"].strip() # 如果文本中检测不到公司名称，但包含12306，则默认为12306 result["company"] = company elif "12306" in text: result["company"] = "12306" # 处理车次格式 if result["trips"]: trips_no = result["trips"].upper() # 清理车次号，只保留字母和数字 trips_no = ''.join(c for c in trips_no if c.isalnum() or c in ['/', '-']) # 验证车次格式 valid_patterns = [re.compile(pattern) for pattern in TrainNERConfig.TRIPS_CONFIG['patterns']] if any(pattern.match(trips_no) for pattern in valid_patterns): result["trips"] = trips_no else: # 尝试修复常见错误 if len(trips_no) >= TrainNERConfig.TRIPS_CONFIG['min_length'] and any(trips_no.startswith(t) for t in TrainNERConfig.TRIPS_CONFIG['train_types']): result["trips"] = trips_no elif trips_no.isdigit() and 1 <= len(trips_no) <= TrainNERConfig.TRIPS_CONFIG['max_length']: result["trips"] = trips_no else: result["trips"] = None # 清理日期格式 if result["date"]: date_str = result["date"] # 保留数字和常见日期分隔符 date_str = ''.join(c for c in date_str if c.isdigit() or c in ['年', '月', '日', '-', '/', '.']) result["date"] = date_str # 清理时间格式 if result["time"]: time_str = result["time"] # 保留数字和常见时间分隔符 time_str = ''.join(c for c in time_str if c.isdigit() or c in [':', '时', '分', '点']) result["time"] = time_str # 处理座位信息 if result["seat"]: seat_str = result["seat"] # 移除可能的额外空格和特殊字符 seat_str = seat_str.replace(" ", "").strip() result["seat"] = seat_str # 处理乘客姓名 if result["name"]: name = result["name"].strip() # 移除可能的标点符号 name = ''.join(c for c in name if c.isalnum() or c in ['*', '·']) result["name"] = name return result except Exception as e: logger.error(f"火车票实体提取失败: {str(e)}") raise # 创建Flask应用 app = Flask(__name__) model_manager = ModelManager() # 添加保存短信到文件的函数 def save_sms_to_file(text: str, category: str = None, confidence: float = None) -> bool: """ 将短信内容保存到本地文件 Args: text: 短信内容 category: 分类结果 confidence: 分类置信度 Returns: bool: 保存成功返回True，否则返回False """ try: # 确保日志目录存在 log_dir = "./sms_logs" if not os.path.exists(log_dir): os.makedirs(log_dir) # 创建基于日期的文件名 today = datetime.datetime.now().strftime("%Y-%m-%d") file_path = os.path.join(log_dir, f"sms_log_{today}.txt") # 获取当前时间 current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") # 准备要写入的内容 category_info = f"分类: {category}, 置信度: {confidence:.4f}" if category and confidence else "未分类" log_content = f"[{current_time}] {category_info}\n{text}\n{'='*50}\n" # 以追加模式写入文件 with open(file_path, 'a', encoding='utf-8') as f: f.write(log_content) return True except Exception as e: logger.error(f"保存短信到文件失败: {str(e)}") return False @app.route("/health", methods=["GET"]) def health_check(): """健康检查接口""" return jsonify({"status": "healthy"}) @app.route("/process-sms", methods=["POST"]) def process_sms(): """处理短信的接口""" try: # 验证请求数据 if not request.is_json: raise BadRequest("请求必须是JSON格式") data = request.get_json() if "content" not in data: raise BadRequest("请求中必须包含'content'字段") text = data["content"] if not isinstance(text, str) or not text.strip(): raise BadRequest("短信内容不能为空") # 保存原始短信内容到文件 save_sms_to_file(text) # 特定短信识别逻辑 - 针对百度通知和招商银行账单 # 识别百度通知 if "百度智能云" in text and "尊敬的用户" in text and "免费额度" in text: logger.info(f"直接识别为百度通知短信: {text[:30]}...") category = "其他" save_sms_to_file(text, category, 1.0) # 记录分类结果 return jsonify({ "status": "success", "data": { "category": category, "details": {} } }) # 识别招商银行账单 if "招商银行" in text and ("账单￥" in text or "账单¥" in text or "还款日" in text): logger.info(f"直接识别为招商银行还款短信: {text[:30]}...") category = "还款" details = model_manager.extract_repayment_entities(text) save_sms_to_file(text, category, 1.0) # 记录分类结果 return jsonify({ "status": "success", "data": { "category": category, "details": details } }) # 处理短信 category, confidence = model_manager.classify_sms(text) # 保存短信内容和分类结果 save_sms_to_file(text, category, confidence) # 如果是明确的业务短信类别，直接进入处理流程 if category in ["快递", "还款", "收入", "航班", "火车票"] and confidence > 0.5: # 对百度通知的特殊处理 if category == "快递" and "百度" in text and "尊敬的用户" in text: logger.info(f"纠正百度通知短信的分类: {text[:30]}...") category = "其他" save_sms_to_file(text, category, confidence) # 更新分类结果 return jsonify({ "status": "success", "data": { "category": category, "details": {} } }) # 对于高置信度的业务分类，直接进入实体提取 if category == "快递": details = model_manager.extract_entities(text) elif category == "还款": details = model_manager.extract_repayment_entities(text) elif category == "收入": details = model_manager.extract_income_entities(text) elif category == "航班": details = model_manager.extract_flight_entities(text) elif category == "火车票": details = model_manager.extract_train_entities(text) logger.info(f"高置信度业务短信: {text[:30]}..., category: {category}, confidence: {confidence:.4f}") return jsonify({ "status": "success", "data": { "category": category, "details": details } }) # 检查是否为营销/广告短信 if model_manager.is_marketing_sms(text): # 如果是营销/广告短信，直接归类为"其他" logger.info(f"检测到营销/广告短信: {text[:30]}...") category = "其他" save_sms_to_file(text, category, confidence) # 更新分类结果 return jsonify({ "status": "success", "data": { "category": category, "details": {} } }) # 检查是否为通知类短信 if model_manager.is_notification_sms(text): # 如果是通知类短信，直接归类为"其他" logger.info(f"检测到通知类短信: {text[:30]}...") category = "其他" save_sms_to_file(text, category, confidence) # 更新分类结果 return jsonify({ "status": "success", "data": { "category": category, "details": {} } }) # 置信度阈值，低于此阈值的分类结果被视为"其他" confidence_threshold = 0.7 if confidence < confidence_threshold: logger.info(f"短信分类置信度低({confidence:.4f})，归类为'其他': {text[:30]}...") category = "其他" save_sms_to_file(text, category, confidence) # 更新分类结果 return jsonify({ "status": "success", "data": { "category": category, "details": {} } }) # 根据分类结果调用对应的实体提取函数 if category == "快递": details = model_manager.extract_entities(text) elif category == "还款": details = model_manager.extract_repayment_entities(text) elif category == "收入": details = model_manager.extract_income_entities(text) elif category == "航班": details = model_manager.extract_flight_entities(text) elif category == "火车票": details = model_manager.extract_train_entities(text) else: details = {} # 记录处理结果 logger.info(f"Successfully processed SMS: {text[:30]}..., category: {category}, confidence: {confidence:.4f}") return jsonify({ "status": "success", "data": { "category": category, "details": details } }) save_sms_to_file except BadRequest as e: logger.warning(f"Invalid request: {str(e)}") return jsonify({ "status": "error", "message": str(e) }), 400 except Exception as e: logger.error(f"Error processing SMS: {str(e)}") return jsonify({ "status": "error", "message": "服务器内部错误" }), 500 if __name__ == "__main__": app.run(host="0.0.0.0", port=5000)