# -*- coding: utf-8 -*-
|
import os
|
import logging
|
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) -> str:
|
"""对短信进行分类"""
|
try:
|
inputs = self.classifier_tokenizer(
|
text,
|
return_tensors="pt",
|
truncation=True,
|
max_length=64
|
)
|
with torch.no_grad():
|
outputs = self.classifier_model(**inputs)
|
pred_id = outputs.logits.argmax().item()
|
return self.classifier_model.config.id2label[pred_id]
|
except Exception as e:
|
logger.error(f"短信分类失败: {str(e)}")
|
raise
|
|
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
|
|
# 在上下文中查找完整金额
|
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,10}(?:\.\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 and value <= 9999999.99: # 设置合理的金额范围
|
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 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, '')
|
|
# 处理金额中的逗号
|
amount_text = amount_text.replace(',', '')
|
|
try:
|
# 转换为浮点数
|
value = float(amount_text)
|
|
# 验证整数位数
|
integer_part = str(int(value))
|
if len(integer_part) <= RepaymentNERConfig.AMOUNT_CONFIG['max_integer_digits']:
|
# 保持原始小数位数
|
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": [],
|
"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_ = "".join(entities["TYPE"]).strip()
|
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_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 amount_candidates:
|
# 按金额大小排序，选择最大的
|
result["amount"] = max(amount_candidates, key=lambda x: x[1])[0]
|
|
# 处理最低还款金额
|
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]
|
|
# 使用正则表达式查找金额模式
|
import re
|
amount_pattern = re.compile(r'(\d{1,10}(?:\.\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
|
|
# 处理收入金额
|
if entities["PICKUP_CODE"]:
|
for amount in entities["PICKUP_CODE"]:
|
cleaned_amount = clean_amount(amount, text)
|
if cleaned_amount:
|
result["amount"] = cleaned_amount
|
break
|
|
# 处理余额
|
if 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()
|
|
@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("短信内容不能为空")
|
|
# 处理短信
|
category = model_manager.classify_sms(text)
|
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]}...")
|
|
return jsonify({
|
"status": "success",
|
"data": {
|
"category": category,
|
"details": details
|
}
|
})
|
|
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)
|
