代码结构优化
删除冗余注释，保留核心功能说明 将Masking类拆分为独立的敏感信息处理模块 使用配置文件管理正则表达式模式和敏感词列表
性能优化
使用多线程处理批量文件 缓存模型加载结果避免重复初始化 添加进度条显示处理进度
错误处理增强
添加文件编码自动检测和异常处理 完善日志记录系统 添加输入验证和类型检查
功能扩展建议
支持更多文件格式如PDF和PPT 添加自定义脱敏规则配置 实现敏感信息检测报告生成
代码规范改进
遵循PEP8命名规范 使用类型注解 统一异常处理方式
安全增强
添加处理前文件备份功能 实现敏感信息替换映射表 支持加密存储脱敏后文件
部署优化
添加Docker支持 配置环境变量管理 实现API服务化封装
测试建议
添加单元测试模块 构建测试数据集 实现自动化测试流程
需要特别注意以下关键点：
- 大模型调用需要处理速率限制
- 文件处理需要保证内存安全
- 敏感词识别要考虑模糊匹配
- 多线程处理需要协调资源竞争
改进后的代码架构建议采用分层设计：
- 应用层：处理CLI和API交互
- 服务层：核心脱敏逻辑
- 数据层：文件读写操作
- 模型层：AI能力封装
数据脱敏处理工具
功能：批量处理文件夹中的Word、Excel、TXT文件，对其中的敏感信息进行脱敏处理
作者：AI助手
日期：2025年
"""
from transformers import AutoModelForCausalLM, AutoTokenizer # 导入Transformers库用于加载大语言模型
import jieba # 中文分词库
import re # 正则表达式库
import pandas as pd # 数据处理库
import argparse # 命令行参数解析库
import json # JSON处理库
import os # 操作系统接口库
import chardet # 文件编码检测库
from vllm import LLM, SamplingParams # vLLM推理加速库
from docx import Document # Word文档处理库
import logging # 日志记录库
# 配置日志输出格式和级别
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
def Masking(text, model, tokenizer):
"""
主要的数据脱敏函数
参数:
text: 需要脱敏的文本内容
model: 加载的大语言模型
tokenizer: 模型对应的分词器
返回:
脱敏后的文本内容
"""
def generating(prompt):
"""
调用大语言模型生成回复的内部函数
参数:
prompt: 输入给模型的提示词
返回:
模型生成的回复文本
"""
# 构建对话消息格式
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": prompt}
]
# 应用聊天模板格式化输入
inputs = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
# 将输入转换为模型可处理的张量格式并移至GPU
model_inputs = tokenizer([inputs], return_tensors="pt").to("cuda")
# 生成回复
generated_ids = model.generate(
model_inputs.input_ids,
max_new_tokens=512, # 最大生成512个新token
pad_token_id=tokenizer.eos_token_id # 设置填充token
)
# 提取新生成的部分（去除输入部分）
generated_ids = [
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
# 解码生成的token为文本
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
return response
def repl4num(match):
"""
手机号码脱敏替换函数
保留前3位和后4位，中间用****替换
参数:
match: 正则匹配对象
返回:
脱敏后的手机号码
"""
num = match.group()
new_num = num[:3] + "****" + num[7:] # 保留前3位和后4位
return new_num
def repl4gonghao(match):
"""
工号脱敏替换函数
保留前3位，其余用*替换
参数:
match: 正则匹配对象
返回:
脱敏后的工号
"""
data = match.group()
new_data = data[:3] + "*" * (len(data) - 3) # 保留前3位，其余用*替换
return new_data
def maskNum(text):
"""
数字信息脱敏函数
处理手机号、工号、中文数字等敏感数字信息
参数:
text: 输入文本
返回:
数字脱敏后的文本
"""
# 匹配11位手机号码的正则表达式
pattern4num = r"\d{11}"
# 匹配工号格式（3个字母+5-7个数字）的正则表达式
pattern4gonghao = r"[a-zA-Z]{3}\d{5,7}"
# 匹配中文数字的正则表达式
pattern4hanzinum = r"[零一幺二三四五六七八九]{3,4}"
# 依次进行各种数字脱敏
text = re.sub(pattern4num, repl4num, text) # 脱敏手机号
text = re.sub(pattern4gonghao, repl4gonghao, text) # 脱敏工号
text = re.sub(pattern4hanzinum, '*', text) # 脱敏中文数字
return text
def maskName(text):
"""
姓名脱敏函数
使用AI模型识别文本中的姓名并进行脱敏
参数:
text: 输入文本
返回:
姓名脱敏后的文本
"""
# 构建识别姓名的提示词
prompt = "下面文本中是否包含姓名，如果包含姓名就直接输出姓名，不包含就输出'没有'：\n" + text
res = generating(prompt) # 调用AI模型识别姓名
# 如果识别到姓名，进行脱敏处理（保留姓氏，名字用[姓名]替换）
if res != '没有':
text = text.replace(res, res[0] + "[姓名]")
return text
def maskAddress(text):
"""
地址脱敏函数
使用AI模型识别文本中的地址信息并进行脱敏
参数:
text: 输入文本
返回:
地址脱敏后的文本
"""
# 构建识别地址的提示词
prompt = "下面文本中是否包含地址信息，如果包含就直接输出地址，不包含就输出'没有'：\n" + text
res = generating(prompt) # 调用AI模型识别地址
# 如果识别到地址，进行脱敏处理
if res != '没有':
res = re.sub(r'[^\w\s]', '', res) # 清理特殊字符
text = text.replace(res, "[地址]") # 替换为[地址]标记
return text
def maskOperator(text):
"""
运营商名称脱敏函数
将文本中的运营商名称替换为通用标记
参数:
text: 输入文本
返回:
运营商脱敏后的文本
"""
# 定义需要脱敏的运营商名称列表
operators = ["中国移动", "联通", "电信"]
# 逐一替换运营商名称
for operator in operators:
text = text.replace(operator, "[运营商]")
return text
# 确保输入是字符串类型
text = str(text)
# 依次执行各种脱敏操作
text = maskNum(text) # 数字脱敏
text = maskName(text) # 姓名脱敏
text = maskAddress(text) # 地址脱敏
text = maskOperator(text) # 运营商脱敏
return text
def data_process(dataset):
"""
文件数据处理函数
根据文件类型读取并解析文件内容
参数:
dataset: 文件路径
返回:
文件内容列表，每个元素是文件中的一行或一个段落
"""
# 检查文件是否存在
if not os.path.exists(dataset):
logging.error(f"文件不存在: {dataset}")
return []
# 处理Word文档（.docx/.doc格式）
if ".docx" in dataset or ".doc" in dataset:
try:
doc = Document(dataset) # 加载Word文档
res = []
# 遍历文档中的每个段落
for para in doc.paragraphs:
if para.text.strip(): # 跳过空段落
res.append(para.text.strip())
return res
except Exception as e:
logging.error(f"无法处理 Word 文件 {dataset}: {str(e)}")
return []
# 处理Excel文件（.xlsx/.xls格式）
elif ".xlsx" in dataset or ".xls" in dataset:
try:
dataset_df = pd.read_excel(dataset) # 读取Excel文件
# 将所有单元格的内容展平为一个列表
return dataset_df.values.flatten().tolist()
except Exception as e:
logging.error(f"无法处理 Excel 文件 {dataset}: {str(e)}")
return []
# 处理文本文件（.txt格式）
elif ".txt" in dataset:
try:
res = []
# 动态检测文件编码
with open(dataset, 'rb') as f: # 以二进制模式读取文件
raw_data = f.read()
encoding_result = chardet.detect(raw_data) # 检测编码
encoding = encoding_result['encoding']
# 使用检测到的编码读取文件
with open(dataset, 'r', encoding=encoding, errors='ignore') as f:
for line in f:
if line.strip(): # 跳过空行
res.append(line.strip())
return res
except Exception as e:
logging.error(f"无法处理 TXT 文件 {dataset}: {str(e)}")
return []
# 不支持的文件格式
return []
def get_all_file_names(folder_path):
"""
递归获取文件夹中所有文件的路径
参数:
folder_path: 文件夹路径
返回:
文件路径列表
"""
try:
file_names = []
# 递归遍历文件夹
for root, dirs, files in os.walk(folder_path):
for file in files:
file_names.append(os.path.join(root, file)) # 获取完整文件路径
# 记录找到的文件
logging.info("文件夹中的所有文件名:")
for file in file_names:
logging.info(file)
return file_names
except FileNotFoundError:
logging.error("错误: 文件夹未找到!")
except Exception as e:
logging.error(f"错误: 发生了一个未知错误: {e}")
return []
def main(args):
"""
主函数
执行批量文件脱敏处理的主要逻辑
参数:
args: 命令行参数对象
"""
# 创建输出文件夹
if not os.path.exists("masked"):
os.makedirs("masked")
# 加载预训练模型和分词器
model = AutoModelForCausalLM.from_pretrained(
args.model_path,
trust_remote_code=True,
device_map='auto' # 自动分配设备（GPU/CPU）
)
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
# 获取待处理的所有文件
datasets = get_all_file_names(args.folder)
# 逐个处理文件
for dataset in datasets:
# 检查文件是否已经处理过
if os.path.exists(f"masked/{os.path.basename(dataset)}"):
logging.info(f"跳过已处理的文件: {dataset}")
continue
logging.info(f'处理文件: {dataset}')
# 读取并解析文件内容
dataset_data = data_process(dataset)
# 如果文件内容为空，跳过该文件
if not dataset_data:
logging.warning(f"跳过处理不好的文件: {dataset}")
continue
# 对文件内容进行脱敏处理
processed_data = []
for data in dataset_data:
masked_data = Masking(data, model=model, tokenizer=tokenizer)
processed_data.append(masked_data)
# 构建输出文件路径
output_file = os.path.join("masked", os.path.basename(dataset))
# 根据文件类型保存脱敏后的内容
if ".docx" in dataset or ".doc" in dataset:
# 保存为Word文档
try:
doc = Document()
for data in processed_data:
doc.add_paragraph(data) # 添加段落
doc.save(output_file)
logging.info(f"Word文件已保存: {output_file}")
except Exception as e:
logging.error(f"保存Word文件失败 {output_file}: {str(e)}")
elif ".xlsx" in dataset or ".xls" in dataset:
# 保存为Excel文件
try:
df = pd.DataFrame(processed_data, columns=['masked_content'])
df.to_excel(output_file, index=False)
logging.info(f"Excel文件已保存: {output_file}")
except Exception as e:
logging.error(f"保存Excel文件失败 {output_file}: {str(e)}")
elif ".txt" in dataset:
# 保存为文本文件
try:
with open(output_file, 'w', encoding='utf-8') as f:
for data in processed_data:
f.write(data + '\n') # 每行一个数据项
logging.info(f"TXT文件已保存: {output_file}")
except Exception as e:
logging.error(f"保存TXT文件失败 {output_file}: {str(e)}")
else:
# 对于其他格式，默认保存为txt文件
try:
output_file = output_file + '.txt'
with open(output_file, 'w', encoding='utf-8') as f:
for data in processed_data:
f.write(str(data) + '\n')
logging.info(f"文件已保存为TXT格式: {output_file}")
except Exception as e:
logging.error(f"保存文件失败 {output_file}: {str(e)}")
if __name__ == "__main__":
# 设置命令行参数解析
parser = argparse.ArgumentParser(description='数据脱敏处理工具')
parser.add_argument('--model_path', type=str, required=True, help='模型路径')
parser.add_argument('--folder', type=str, required=True, help='要处理的文件夹路径')
args = parser.parse_args()
# 执行主函数
main(args)
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
