self-tools
/
bert_base
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

1

This commit is contained in:
Rzy 2026-08-11 14:27:21 +08:00
parent 0d5060489a
commit de92a67288
12 changed files with 2338 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
.env Normal file
View File

@ -0,0 +1,22 @@
# ==================== ModelScope配置 ====================
MODELSCOPE_TOKEN=ms-00353db2-f626-47b9-8c59-d535824fd7fb
MODEL_CACHE_DIR=/app/model_cache
# ==================== 税务分类模型 ====================
TAX_MODEL_ID=circles1/tax_rate
TAX_MODEL_DIR=/app/services/tax/model
TAX_PORT=5003
TAX_SERVICE_NAME=tax_classifier
# ==================== 订单地址模型 ====================
ADDRESS_MODEL_ID=circles1/order_address
ADDRESS_MODEL_DIR=/app/services/bert_address/model
ADDRESS_PORT=5002
ADDRESS_SERVICE_NAME=address_classifier
# ==================== 运行配置 ====================
DEVICE=cuda
MAX_WORKERS=4
BATCH_SIZE=32

23
Dockerfile Normal file
View File

@ -0,0 +1,23 @@
FROM registry.cn-chengdu.aliyuncs.com/go_ls/bert_base:latest
WORKDIR /app
COPY services/ ./services/
COPY scripts/ ./scripts/
COPY .env .env
# 生成 supervisor 配置
RUN python scripts/generate_supervisor_conf.py
# 创建日志目录并设置权限
RUN mkdir -p /var/log/supervisor /app/logs && \
chown -R 1000:1000 /var/log/supervisor /app/logs
# 创建非root用户
RUN useradd -m -u 1000 appuser && chown -R appuser:appuser /app
USER appuser
# 启动命令：先下载模型，再启动 supervisor
CMD ["sh", "-c", "python scripts/download_models.py && supervisord -c ./supervisor.conf"]

23
base/Dockerfile Normal file
View File

@ -0,0 +1,23 @@
FROM python:3.12-slim
# 设置工作目录
WORKDIR /app
# 设置环境变量
ENV PYTHONUNBUFFERED=1 \
PYTHONDONTWRITEBYTECODE=1 \
PIP_NO_CACHE_DIR=1 \
PIP_DISABLE_PIP_VERSION_CHECK=1 \
DEBIAN_FRONTEND=noninteractive
# 安装系统依赖
RUN apt-get update && apt-get install -y \
gcc \
g++ \
supervisor \
&& rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

9
base/requirements.txt Normal file
View File

@ -0,0 +1,9 @@
Flask==3.1.3
torch==2.13.0
transformers==5.14.1
joblib==1.5.3
gunicorn==26.0.0
requests==2.34.2
modelscope==1.39.1
load_dotenv==0.1.0
scikit-learn

2
deploy.sh Normal file
View File

@ -0,0 +1,2 @@
# 3. 首次启动（会下载模型到本地）
docker-compose up -d

39
docker-compose.yml Normal file
View File

@ -0,0 +1,39 @@
services:
bert-all-services:
build:
context: .
dockerfile: Dockerfile
container_name: bert-all-services
ports:
- "5002-5010:5002-5010"
env_file:
- .env
environment:
- DEVICE=${DEVICE:-cuda}
- MAX_WORKERS=${MAX_WORKERS:-4}
- BATCH_SIZE=${BATCH_SIZE:-32}
- NVIDIA_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES:-all}
- NVIDIA_DRIVER_CAPABILITIES=${NVIDIA_DRIVER_CAPABILITIES:-all}
- MODEL_CACHE_DIR=/app/model_cache
# 传递代理环境变量（如果构建时使用了代理）
- HTTP_PROXY=${HTTP_PROXY:-}
- HTTPS_PROXY=${HTTPS_PROXY:-}
- http_proxy=${http_proxy:-}
- https_proxy=${https_proxy:-}
- NO_PROXY=${NO_PROXY:-}
- no_proxy=${no_proxy:-}
volumes:
- ./logs:/app/logs
- ./models:/app/models
- ./model_cache:/app/model_cache
restart: unless-stopped
deploy:
resources:
limits:
memory: 8G
reservations:
memory: 4G
devices:
- driver: nvidia
count: all
capabilities: [gpu]

276
scripts/download_models.py Normal file
View File

@ -0,0 +1,276 @@
import os
import sys
import shutil
from modelscope.hub.snapshot_download import snapshot_download
from dotenv import load_dotenv
load_dotenv()
# 使用持久化的缓存目录（挂载到宿主机）
CACHE_DIR = os.getenv('MODEL_CACHE_DIR', '/app/model_cache')
os.makedirs(CACHE_DIR, exist_ok=True)
def discover_models():
"""从环境变量中发现所有模型"""
models = {}
for key, value in os.environ.items():
if key.endswith('_MODEL_ID'):
prefix = key[:-9]
service_name = os.getenv(f'{prefix}_SERVICE_NAME', f'{prefix.lower()}_classifier')
service_folder = service_name.replace('_classifier', '')
models[prefix.lower()] = {
'prefix': prefix,
'model_id': value,
'port': os.getenv(f'{prefix}_PORT', '5002'),
'service_name': service_name,
'model_dir': os.getenv(f'{prefix}_MODEL_DIR', f'/app/services/bert_{service_folder}/model'),
'service_folder': service_folder,
'cache_dir': os.path.join(CACHE_DIR, prefix.lower())
}
return models
def check_model_complete(model_dir):
"""检查模型是否完整"""
if not os.path.exists(model_dir):
return False
try:
files = os.listdir(model_dir)
if not files:
return False
has_weights = any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in files)
has_config = any(f == 'config.json' for f in files)
if has_weights and has_config:
return True
return False
except Exception:
return False
def find_model_files(cache_dir):
"""在下载的缓存目录中查找实际的模型文件"""
# 检查是否直接有模型文件
direct_files = [f for f in os.listdir(cache_dir) if
not f.startswith('.') and os.path.isfile(os.path.join(cache_dir, f))]
if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in direct_files):
return cache_dir
# 检查 models/ 子目录（ModelScope 的标准结构）
models_dir = os.path.join(cache_dir, 'models')
if os.path.exists(models_dir):
for root, dirs, files in os.walk(models_dir):
if 'snapshots' in root:
for f in files:
if f.endswith(('.safetensors', '.bin', '.pt', '.pth')):
return root
if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in files):
return root
# 检查是否存在单层子目录
subdirs = [d for d in os.listdir(cache_dir) if os.path.isdir(os.path.join(cache_dir, d)) and not d.startswith('.')]
for subdir in subdirs:
subpath = os.path.join(cache_dir, subdir)
if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in os.listdir(subpath)):
return subpath
# 递归查找任何包含模型文件的目录
for root, dirs, files in os.walk(cache_dir):
if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in files):
return root
return None
def copy_model_from_cache(cache_dir, target_dir):
"""从缓存目录复制模型到目标目录"""
src_dir = find_model_files(cache_dir)
if src_dir is None:
print(f" ✗ Could not find model files in cache: {cache_dir}")
return False
print(f" ✓ Found cached model files in: {src_dir}")
os.makedirs(target_dir, exist_ok=True)
copied_count = 0
for file in os.listdir(src_dir):
src_file = os.path.join(src_dir, file)
if os.path.isfile(src_file):
dest_file = os.path.join(target_dir, file)
shutil.copy2(src_file, dest_file)
copied_count += 1
print(f" ✓ Copied {copied_count} files to {target_dir}")
return copied_count > 0
def download_model_with_fallback(name, config):
"""下载单个模型，带备选方案"""
print("=" * 50)
print(f"Processing {name} model...")
print(f"Model ID: {config['model_id']}")
print(f"Target: {config['model_dir']}")
print(f"Cache: {config['cache_dir']}")
print("=" * 50)
try:
os.makedirs(config['model_dir'], exist_ok=True)
os.makedirs(config['cache_dir'], exist_ok=True)
# 1. 检查目标目录是否已有完整模型
if check_model_complete(config['model_dir']):
print(f"✅ Model already exists in target, skipping")
return True
# 2. 检查缓存目录是否有模型
if check_model_complete(config['cache_dir']):
print(f"✅ Found complete model in local cache")
print(f" Copying from cache to target...")
if copy_model_from_cache(config['cache_dir'], config['model_dir']):
if check_model_complete(config['model_dir']):
print(f"✅ Model copied successfully from cache")
return True
else:
print(f"⚠️ Model copied but appears incomplete, will re-download")
# 3. 缓存中没有，需要下载
print(f"📥 Downloading model from ModelScope...")
download_success = False
# 尝试方式1：直接下载
try:
snapshot_download(
model_id=config['model_id'],
cache_dir=config['cache_dir'],
revision="master",
ignore_file_pattern=[".git", ".gitattributes"]
)
print(f"✅ Download completed")
download_success = True
except Exception as e:
print(f"⚠️ Primary download failed: {e}")
# 尝试方式2：使用 HubApi（如果有 token）
token = os.getenv('MODELSCOPE_TOKEN', '')
if token:
print(" Trying alternative download with HubApi...")
try:
from modelscope.hub.api import HubApi
api = HubApi()
api.login(token=token)
snapshot_download(
model_id=config['model_id'],
cache_dir=config['cache_dir'],
revision="master",
ignore_file_pattern=[".git", ".gitattributes"]
)
print(f"✅ Alternative download completed")
download_success = True
except Exception as alt_error:
print(f"✗ Alternative download also failed: {alt_error}")
raise
else:
print("✗ No MODELSCOPE_TOKEN set, cannot use alternative method")
raise
if not download_success:
return False
# 验证并复制
if not check_model_complete(config['cache_dir']):
print(f"⚠️ Downloaded model appears incomplete in cache")
print(f" Attempting to copy anyway...")
print(f"Copying from cache to target...")
if copy_model_from_cache(config['cache_dir'], config['model_dir']):
if check_model_complete(config['model_dir']):
print(f"✅ Model deployed successfully!")
return True
else:
print(f"⚠️ Model deployed but appears incomplete")
return True
else:
print(f"✗ Failed to copy model to target")
return False
except Exception as e:
print(f"✗ Failed to process {name} model: {e}")
import traceback
traceback.print_exc()
return False
def main():
print("=" * 50)
print("Model Download Script")
print("=" * 50)
print(f"Cache directory: {CACHE_DIR}")
force_download = os.getenv('FORCE_MODEL_DOWNLOAD', 'false').lower() == 'true'
if force_download:
print("⚠️ FORCE_MODEL_DOWNLOAD=true - will re-download all models")
print("\nDiscovering models from .env...")
models = discover_models()
if not models:
print("⚠️ No models found in .env file!")
sys.exit(1)
print(f"\nFound {len(models)} model(s):")
all_exist = True
for name, config in models.items():
target_exists = check_model_complete(config['model_dir'])
cache_exists = check_model_complete(config['cache_dir'])
print(f" {name}: {config['model_id']}")
print(f" Target: {'✅' if target_exists else '❌'} {config['model_dir']}")
print(f" Cache: {'✅' if cache_exists else '❌'} {config['cache_dir']}")
if not target_exists and not cache_exists:
all_exist = False
print("")
if all_exist and not force_download:
print("✅ All models exist!")
print(" Skipping downloads.")
sys.exit(0)
if force_download:
print("⚠️ Force download enabled, clearing caches...")
for name, config in models.items():
if os.path.exists(config['cache_dir']):
print(f" Clearing cache for {name}")
shutil.rmtree(config['cache_dir'])
os.makedirs(config['cache_dir'], exist_ok=True)
print("\n" + "=" * 50)
print("Processing models...")
print("=" * 50)
success_count = 0
for name, config in models.items():
if download_model_with_fallback(name, config):
success_count += 1
print("")
print("=" * 50)
print(f"Processing completed: {success_count}/{len(models)} models")
print("=" * 50)
if success_count < len(models):
sys.exit(1)
if __name__ == "__main__":
main()

104
scripts/generate_supervisor_conf.py Normal file
View File

@ -0,0 +1,104 @@
import os
import sys
from dotenv import load_dotenv
load_dotenv()
def discover_models():
"""从环境变量中发现所有模型"""
models = {}
for key, value in os.environ.items():
if key.endswith('_MODEL_ID'):
prefix = key[:-9]
models[prefix.lower()] = {
'prefix': prefix,
'model_id': value,
'port': os.getenv(f'{prefix}_PORT', '5002'),
'service_name': os.getenv(f'{prefix}_SERVICE_NAME', f'{prefix.lower()}_classifier'),
'model_dir': os.getenv(f'{prefix}_MODEL_DIR', f'/app/services/{prefix.lower()}/model')
}
return models
def generate_supervisor_conf(models):
"""生成supervisor配置文件"""
conf = []
# supervisor全局配置
conf.append("""[supervisord]
nodaemon=true
logfile=/var/log/supervisor/supervisord.log
pidfile=/var/run/supervisord.pid
childlogdir=/var/log/supervisor
""")
# 为每个模型生成program配置
programs = []
for name, config in models.items():
program_name = config['service_name']
programs.append(program_name)
conf.append(f"""[program:{program_name}]
command=gunicorn --bind 0.0.0.0:{config['port']} --workers %(ENV_MAX_WORKERS)s --threads 2 --timeout 120 services.{name}.app:app
directory=/app
autostart=true
autorestart=true
startretries=3
stdout_logfile=/app/logs/{program_name}.log
stdout_logfile_maxbytes=50MB
stderr_logfile=/app/logs/{program_name}_error.log
stderr_logfile_maxbytes=50MB
environment=SERVICE_NAME="{program_name}",SERVICE_PORT="{config['port']}",MODEL_ID="{config['model_id']}"
user=appuser
""")
# 生成group配置
conf.append(f"""[group:bert_services]
programs={','.join(programs)}
""")
return ''.join(conf)
def main():
supervisor_conf_path = './supervisor.conf'
# 检查文件是否已存在
if os.path.exists(supervisor_conf_path):
print("=" * 50)
print("supervisor.conf already exists, skipping generation...")
print(f"File: {os.path.abspath(supervisor_conf_path)}")
print("=" * 50)
sys.exit(0)
print("=" * 50)
print("Generating supervisor.conf dynamically...")
print("=" * 50)
models = discover_models()
if not models:
print("⚠️ No models found, generating empty config")
sys.exit(1)
print(f"Found {len(models)} model(s):")
for name, config in models.items():
print(f" - {name}: {config['model_id']} (port {config['port']})")
# 生成配置
conf_content = generate_supervisor_conf(models)
# 写入文件
with open(supervisor_conf_path, 'w') as f:
f.write(conf_content)
print(f"\n✓ supervisor.conf generated successfully!")
print(f" File: {os.path.abspath(supervisor_conf_path)}")
print("=" * 50)
if __name__ == "__main__":
main()

349
services/address/app.py Normal file
View File

@ -0,0 +1,349 @@
import os
import sys
from flask import Flask, request, jsonify
import torch
from transformers import BertTokenizer, BertForSequenceClassification
import joblib
import re
from functools import lru_cache
from concurrent.futures import ThreadPoolExecutor
from typing import List, Dict
import threading
import logging
import atexit
from datetime import datetime
import time
# ==================== 从环境变量读取配置 ====================
SERVICE_NAME = os.environ.get('SERVICE_NAME', 'address_classifier')
SERVICE_PORT = int(os.environ.get('SERVICE_PORT', 5003))
MODEL_ID = os.environ.get('MODEL_ID', 'circles1/address_model')
MODEL_DIR = os.environ.get('MODEL_DIR', '/app/services/bert_address/model')
# 模型文件路径
LABEL_ENCODER_PATH = os.path.join(MODEL_DIR, 'label_encoder.pkl')
# 如果上面路径不存在，尝试其他可能的文件名
if not os.path.exists(LABEL_ENCODER_PATH):
import glob
encoder_files = glob.glob(os.path.join(MODEL_DIR, 'label_encoder*.pkl'))
if encoder_files:
LABEL_ENCODER_PATH = encoder_files[0]
# ==================== 配置参数 ====================
MAX_LENGTH = 512
DEVICE = os.environ.get('DEVICE', 'cuda' if torch.cuda.is_available() else 'cpu')
# 如果设备设置为cuda但不可用，回退到cpu
if DEVICE == 'cuda' and not torch.cuda.is_available():
DEVICE = 'cpu'
MAX_WORKERS = int(os.environ.get('MAX_WORKERS', 4))
CACHE_SIZE = int(os.environ.get('CACHE_SIZE', 2000))
TOKEN_CACHE_SIZE = int(os.environ.get('TOKEN_CACHE_SIZE', 1000))
# ==================== 日志配置 ====================
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# 打印启动信息
logger.info("=" * 50)
logger.info(f"Starting service: {SERVICE_NAME}")
logger.info(f"Service port: {SERVICE_PORT}")
logger.info(f"Model ID: {MODEL_ID}")
logger.info(f"Model directory: {MODEL_DIR}")
logger.info(f"Label encoder: {LABEL_ENCODER_PATH}")
logger.info(f"Device: {DEVICE}")
logger.info(f"Max workers: {MAX_WORKERS}")
logger.info("=" * 50)
app = Flask(__name__)
app.config['JSON_AS_ASCII'] = False
# 全局变量锁
model_lock = threading.Lock()
class Predictor:
"""预测器类，用于管理模型和分词器的生命周期"""
_instance = None
def __new__(cls):
if cls._instance is None:
with model_lock:
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance._initialized = False
return cls._instance
def __init__(self):
if self._initialized:
return
logger.info(f"Initializing Predictor on device: {DEVICE}")
start_time = time.time()
try:
# 检查模型目录是否存在
if not os.path.exists(MODEL_DIR):
raise FileNotFoundError(f"Model directory not found: {MODEL_DIR}")
# 检查label_encoder是否存在
if not os.path.exists(LABEL_ENCODER_PATH):
raise FileNotFoundError(f"Label encoder not found: {LABEL_ENCODER_PATH}")
# 加载模型和分词器
logger.info(f"Loading tokenizer and model from {MODEL_DIR}...")
self.tokenizer = BertTokenizer.from_pretrained(MODEL_DIR)
self.model = BertForSequenceClassification.from_pretrained(MODEL_DIR).to(DEVICE)
self.model.eval()
logger.info("Model loaded successfully")
# 加载标签映射器
logger.info(f"Loading label encoder from {LABEL_ENCODER_PATH}...")
self.label_encoder = joblib.load(LABEL_ENCODER_PATH)
self.label_map = {i: label for i, label in enumerate(self.label_encoder.classes_)}
logger.info(f"Loaded {len(self.label_map)} labels: {list(self.label_map.values())}")
# 创建线程池
self.executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)
logger.info(f"ThreadPoolExecutor created with {MAX_WORKERS} workers")
# 注册关闭钩子
atexit.register(self._shutdown)
# 模型预热
self._warmup()
self._initialized = True
elapsed_time = time.time() - start_time
logger.info(f"Predictor initialized successfully in {elapsed_time:.2f}s")
except Exception as e:
logger.error(f"Failed to initialize Predictor: {e}")
raise
def _shutdown(self):
"""清理资源"""
logger.info("Shutting down executor...")
self.executor.shutdown(wait=True)
logger.info("Executor shutdown complete")
def _warmup(self):
"""模型预热，确保CUDA初始化"""
logger.info("Warming up model...")
try:
test_text = "测试地址"
self.predict_single(test_text)
logger.info("Model warmup completed")
except Exception as e:
logger.warning(f"Model warmup failed: {e}")
@lru_cache(maxsize=CACHE_SIZE)
def clean_text(self, text: str) -> str:
"""文本清洗函数，带缓存"""
if not isinstance(text, str):
return ""
# 只保留中文字符
cleaned_text = re.sub(r'[^\u4e00-\u9fa5]', '', text)
return cleaned_text.strip()
def predict_single(self, text: str) -> Dict:
"""单个文本预测"""
if not text or not isinstance(text, str):
return {"address": "", "confidence": 0.0, "error": "Invalid input"}
try:
cleaned_text = self.clean_text(text)
if not cleaned_text:
logger.warning(f"Empty text after cleaning: {text[:50]}...")
return {"address": "", "confidence": 0.0}
inputs = self.tokenizer(
cleaned_text,
return_tensors="pt",
truncation=True,
padding=True,
max_length=MAX_LENGTH
).to(DEVICE)
with torch.no_grad():
outputs = self.model(**inputs)
probs = torch.softmax(outputs.logits, dim=1).cpu()
top_prob, top_idx = torch.topk(probs, k=1)
return {
"address": self.label_map[top_idx.item()],
"confidence": round(top_prob.item(), 4)
}
except torch.cuda.OutOfMemoryError as e:
logger.error(f"CUDA OOM error: {e}")
torch.cuda.empty_cache()
return {"address": "", "confidence": 0.0, "error": "GPU memory exhausted"}
except Exception as e:
logger.error(f"Prediction error for text '{text[:50]}...': {e}")
return {"address": "", "confidence": 0.0, "error": str(e)}
def batch_predict(self, texts: List[str]) -> List[Dict]:
"""批量预测（并发处理）"""
if not texts:
return []
logger.info(f"Processing {len(texts)} texts")
start_time = time.time()
# 使用线程池并发处理
futures = [self.executor.submit(self.predict_single, text) for text in texts]
results = [future.result() for future in futures]
elapsed_time = time.time() - start_time
logger.info(f"Batch prediction completed in {elapsed_time:.2f}s for {len(texts)} texts")
return results
def get_stats(self) -> Dict:
"""获取预测器统计信息"""
return {
"service_name": SERVICE_NAME,
"model_id": MODEL_ID,
"model_dir": MODEL_DIR,
"device": DEVICE,
"num_labels": len(self.label_map),
"labels": list(self.label_map.values()),
"max_length": MAX_LENGTH,
"max_workers": MAX_WORKERS
}
# 初始化预测器
try:
predictor = Predictor()
logger.info("Predictor instance created successfully")
except Exception as e:
logger.error(f"Failed to create predictor instance: {e}")
raise
@app.route('/predict', methods=['POST'])
def predict():
"""预测接口"""
try:
data = request.get_json()
if not data:
return jsonify({
"status": "error",
"error": "Invalid request, JSON body required"
}), 400
if 'address' not in data:
return jsonify({
"status": "error",
"error": "Invalid request, 'address' is required"
}), 400
addresses = data['address']
if not isinstance(addresses, list):
addresses = [addresses]
# 限制最大处理数量
max_items = data.get('max_items', 1000)
if len(addresses) > max_items:
logger.warning(f"Too many items: {len(addresses)} > {max_items}, truncating")
addresses = addresses[:max_items]
# 去重（可选）
if data.get('deduplicate', False):
addresses = list(dict.fromkeys(addresses))
logger.info(f"After deduplication: {len(addresses)} unique texts")
# 批量预测
results = predictor.batch_predict(addresses)
return jsonify({
"status": "success",
"predictions": results,
"metadata": {
"total": len(results),
"timestamp": datetime.now().isoformat(),
"service": SERVICE_NAME,
"device": DEVICE
}
})
except Exception as e:
logger.error(f"Prediction endpoint error: {e}", exc_info=True)
return jsonify({
"status": "error",
"error": f"Internal server error: {str(e)}"
}), 500
@app.route('/health', methods=['GET'])
def health_check():
"""健康检查接口"""
try:
stats = predictor.get_stats()
return jsonify({
"status": "healthy",
"timestamp": datetime.now().isoformat(),
"stats": stats
})
except Exception as e:
return jsonify({
"status": "unhealthy",
"error": str(e)
}), 500
@app.route('/stats', methods=['GET'])
def get_stats():
"""获取预测器统计信息"""
try:
stats = predictor.get_stats()
return jsonify({
"status": "success",
"stats": stats
})
except Exception as e:
return jsonify({
"status": "error",
"error": str(e)
}), 500
@app.errorhandler(404)
def not_found(error):
return jsonify({
"status": "error",
"error": "Endpoint not found"
}), 404
@app.errorhandler(500)
def internal_error(error):
return jsonify({
"status": "error",
"error": "Internal server error"
}), 500
if __name__ == '__main__':
logger.info(f"Starting Flask application: {SERVICE_NAME}")
logger.info(f"Model path: {MODEL_DIR}")
logger.info(f"Device: {DEVICE}")
logger.info(f"Port: {SERVICE_PORT}")
app.run(
host='0.0.0.0',
port=SERVICE_PORT,
threaded=True,
debug=False
)
else:
# 兼容 WSGI 标准（如 Gunicorn）
application = app
logger.info(f"Application loaded for WSGI server: {SERVICE_NAME}")

215
services/address/train.py Normal file
View File

@ -0,0 +1,215 @@
import pandas as pd
from sklearn.model_selection import train_test_split
from transformers import (
BertTokenizer,
BertForSequenceClassification,
Trainer,
TrainingArguments,
EarlyStoppingCallback
)
import torch
from torch.utils.data import Dataset
from tqdm import tqdm
import warnings
import re # 用于正则表达式清洗
from sklearn.preprocessing import LabelEncoder
import joblib
# 1. 参数配置（集中管理）
class Config:
MODEL_NAME = "bert-base-chinese"
MAX_LENGTH = 64
BATCH_SIZE = 32
NUM_EPOCHS = 5
LEARNING_RATE = 2e-5
WARMUP_STEPS = 500
WEIGHT_DECAY = 0.01
FP16 = torch.cuda.is_available()
OUTPUT_DIR = "./results/single_level"
LOG_DIR = "./logs"
SAVE_DIR = "./saved_model/single_level"
DEVICE = "cuda" if FP16 else "cpu"
# 2. 数据加载与预处理（添加异常处理和日志）
def load_data(file_path):
try:
df = pd.read_csv(file_path)
assert {'sentence', 'label'}.issubset(df.columns), "数据必须包含'sentence'和'label'列"
print(f"✅ 数据加载成功 | 样本量: {len(df)} | 分类数: {df['label'].nunique()}")
return df
except Exception as e:
warnings.warn(f"❌ 数据加载失败: {str(e)}")
raise
# 新增：数据清洗函数 - 只保留中文字符
def clean_chinese_text(text):
"""
清洗文本，只保留中文字符
"""
if not isinstance(text, str):
return ""
# 使用正则表达式匹配所有中文字符（包括中文标点符号）[^\u4e00-\u9fa5\u3000-\u303f\uff00-\uffef]
# 如果需要更严格的只保留汉字，可以使用：[\u4e00-\u9fa5]
cleaned_text = re.sub(r'[^\u4e00-\u9fa5]', '', text)
return cleaned_text.strip()
# 3. 优化Dataset（添加内存缓存和批处理支持）
class TextDataset(Dataset):
def __init__(self, dataframe, tokenizer, text_col="sentence", label_col="label"):
self.data = dataframe
self.tokenizer = tokenizer
self.text_col = text_col
self.label_col = label_col
# 预计算编码（空间换时间）
self.encodings = tokenizer(
dataframe[text_col].tolist(),
max_length=Config.MAX_LENGTH,
padding="max_length",
truncation=True,
return_tensors="pt"
)
self.labels = torch.tensor(dataframe[label_col].values, dtype=torch.long)
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return {
"input_ids": self.encodings["input_ids"][idx],
"attention_mask": self.encodings["attention_mask"][idx],
"labels": self.labels[idx]
}
# 4. 模型初始化（添加设备移动）
def init_model(num_labels):
tokenizer = BertTokenizer.from_pretrained(Config.MODEL_NAME)
model = BertForSequenceClassification.from_pretrained(
Config.MODEL_NAME,
num_labels=num_labels,
ignore_mismatched_sizes=True # 可选
).to(Config.DEVICE)
return tokenizer, model
# 5. 训练配置（添加早停和梯度累积）
def get_training_args():
return TrainingArguments(
output_dir=Config.OUTPUT_DIR,
num_train_epochs=Config.NUM_EPOCHS,
per_device_train_batch_size=Config.BATCH_SIZE,
per_device_eval_batch_size=Config.BATCH_SIZE * 2, # 评估时可用更大batch
learning_rate=Config.LEARNING_RATE,
warmup_steps=Config.WARMUP_STEPS,
weight_decay=Config.WEIGHT_DECAY,
logging_dir=Config.LOG_DIR,
logging_steps=10,
eval_strategy="steps",
eval_steps=100,
save_strategy="steps",
save_steps=200,
load_best_model_at_end=True,
metric_for_best_model="eval_loss",
greater_is_better=False,
fp16=Config.FP16,
gradient_accumulation_steps=2, # 模拟更大batch
report_to="none", # 禁用wandb等报告
seed=42
)
# 6. 优化推理函数（添加批处理支持）
@torch.no_grad()
def batch_predict(texts, model, tokenizer, label_map, top_k=1, batch_size=16):
model.eval()
all_results = []
for i in tqdm(range(0, len(texts), batch_size), desc="预测中"):
batch = texts[i:i + batch_size]
inputs = tokenizer(
batch,
return_tensors="pt",
truncation=True,
padding=True,
max_length=Config.MAX_LENGTH
).to(Config.DEVICE)
outputs = model(**inputs)
probs = torch.softmax(outputs.logits, dim=1).cpu()
for prob in probs:
top_probs, top_indices = torch.topk(prob, k=top_k)
all_results.extend([
{
"category": label_map[idx.item()],
"confidence": prob.item()
}
for prob, idx in zip(top_probs, top_indices)
])
return all_results[:len(texts)] # 处理非整除情况
# 主流程
if __name__ == "__main__":
# 1. 加载数据
df = load_data("order_address.csv")
# 2. 数据清洗 - 只保留中文
print("🧼 开始清洗文本数据...")
df['sentence'] = df['sentence'].apply(clean_chinese_text)
df = df[df['sentence'].str.len() > 0].reset_index(drop=True)
print(f"✅ 数据清洗完成 | 剩余样本量: {len(df)}")
# 3. 处理中文标签：映射为数值ID，并保存映射关系
print("🏷️ 处理中文标签...")
label_encoder = LabelEncoder()
df['label_id'] = label_encoder.fit_transform(df['label']) # 中文标签 → 数值ID
label_map = {i: label for i, label in enumerate(label_encoder.classes_)} # 数值ID → 中文标签
print(f"标签映射示例: {label_map}")
# 保存标签映射器（供推理时使用）
joblib.dump(label_encoder, "cate/label_encoder.pkl")
print(f"✅ 标签映射完成 | 类别数: {len(label_map)}")
# 4. 划分数据集（使用 label_id 列）
train_df, test_df = train_test_split(
df, test_size=0.2, random_state=42, stratify=df["label_id"] # 注意这里用 label_id
)
# 5. 初始化模型（使用数值标签的数量）
num_labels = len(label_map)
tokenizer, model = init_model(num_labels)
# 6. 准备数据集（使用 label_id 列）
train_dataset = TextDataset(train_df, tokenizer, label_col="label_id") # 指定 label_col
test_dataset = TextDataset(test_df, tokenizer, label_col="label_id")
# 7. 训练配置（保持不变）
training_args = get_training_args()
# 8. 训练器（保持不变）
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=test_dataset,
callbacks=[EarlyStoppingCallback(early_stopping_patience=3)]
)
# 9. 训练和保存（保持不变）
trainer.train()
model.save_pretrained(Config.SAVE_DIR)
tokenizer.save_pretrained(Config.SAVE_DIR)
# 12. 测试推理
test_samples = ["山东省济南市莱芜区碧桂园天樾422502", "广东省广州市花都区狮岭镇山前旅游大道18号机车检修段", "江苏省苏州市吴中区吴中区木渎镇枫瑞路85号诺德·长枫雅苑北区10栋-303"]
# 先清洗测试样本
cleaned_samples = [clean_chinese_text(s) for s in test_samples]
predictions = batch_predict(cleaned_samples, model, tokenizer, label_map)
for sample, pred in zip(test_samples, predictions):
print(
f"输入: {sample}\n清洗后: {clean_chinese_text(sample)}\n预测: {pred['category']} (置信度: {pred['confidence']:.2f})\n")

451
services/tax/app.py Normal file
View File

@ -0,0 +1,451 @@
import os
import sys
from flask import Flask, request, jsonify
import torch
from transformers import BertTokenizer, BertForSequenceClassification
import joblib
import re
from functools import lru_cache
from concurrent.futures import ThreadPoolExecutor
from typing import List, Dict
import threading
import logging
import atexit
from datetime import datetime
import time
# ==================== 从环境变量读取配置 ====================
SERVICE_NAME = os.environ.get('SERVICE_NAME', 'tax_classifier')
SERVICE_PORT = int(os.environ.get('SERVICE_PORT', 5002))
MODEL_ID = os.environ.get('MODEL_ID', 'circles1/tax_rate')
MODEL_DIR = os.environ.get('MODEL_DIR', '/app/services/tax_classifier/model')
# 兼容旧版路径（如果MODEL_DIR未设置，尝试从SERVICE_NAME推断）
if not os.environ.get('MODEL_DIR'):
# 从SERVICE_NAME推断模型目录
service_folder = SERVICE_NAME.replace('_classifier', '')
MODEL_DIR = f'/app/services/{service_folder}/model'
# 模型文件路径
LABEL_ENCODER_PATH = os.path.join(MODEL_DIR, 'label_encoder_roberta_large.pkl')
# 如果上面路径不存在，尝试在模型根目录查找
if not os.path.exists(LABEL_ENCODER_PATH):
LABEL_ENCODER_PATH = os.path.join(MODEL_DIR, 'label_encoder.pkl')
# ==================== 配置参数 ====================
MAX_LENGTH = 512
DEVICE = os.environ.get('DEVICE', 'cuda' if torch.cuda.is_available() else 'cpu')
# 如果设备设置为cuda但不可用，回退到cpu
if DEVICE == 'cuda' and not torch.cuda.is_available():
DEVICE = 'cpu'
print(f"⚠️ CUDA not available, falling back to {DEVICE}")
BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 32))
MAX_WORKERS = int(os.environ.get('MAX_WORKERS', 4))
CACHE_SIZE = 2000
TOKEN_CACHE_SIZE = 1000
# ==================== 日志配置 ====================
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# 打印启动信息
logger.info("=" * 50)
logger.info(f"Starting service: {SERVICE_NAME}")
logger.info(f"Service port: {SERVICE_PORT}")
logger.info(f"Model ID: {MODEL_ID}")
logger.info(f"Model directory: {MODEL_DIR}")
logger.info(f"Device: {DEVICE}")
logger.info(f"Batch size: {BATCH_SIZE}")
logger.info(f"Max workers: {MAX_WORKERS}")
logger.info("=" * 50)
app = Flask(__name__)
app.config['JSON_AS_ASCII'] = False
# 全局变量锁
model_lock = threading.Lock()
class Predictor:
"""预测器类，用于管理模型和分词器的生命周期"""
_instance = None
def __new__(cls):
if cls._instance is None:
with model_lock:
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance._initialized = False
return cls._instance
def __init__(self):
if self._initialized:
return
global LABEL_ENCODER_PATH # 关键修复
logger.info(f"Initializing Predictor on device: {DEVICE}")
start_time = time.time()
try:
# 检查模型目录是否存在
if not os.path.exists(MODEL_DIR):
raise FileNotFoundError(f"Model directory not found: {MODEL_DIR}")
# 检查label_encoder是否存在
if not os.path.exists(LABEL_ENCODER_PATH):
# 尝试在目录中查找label_encoder文件
import glob
encoder_files = glob.glob(os.path.join(MODEL_DIR, 'label_encoder*.pkl'))
if encoder_files:
actual_encoder_path = encoder_files[0]
logger.info(f"Found label_encoder at: {actual_encoder_path}")
LABEL_ENCODER_PATH = actual_encoder_path
else:
raise FileNotFoundError(f"Label encoder not found in {MODEL_DIR}")
# 加载模型和分词器
logger.info(f"Loading tokenizer and model from {MODEL_DIR}...")
self.tokenizer = BertTokenizer.from_pretrained(MODEL_DIR)
self.model = BertForSequenceClassification.from_pretrained(MODEL_DIR).to(DEVICE)
self.model.eval()
logger.info(f"Model loaded successfully")
# 加载标签映射器
logger.info(f"Loading label encoder from {LABEL_ENCODER_PATH}...")
self.label_encoder = joblib.load(LABEL_ENCODER_PATH)
self.num_labels = len(self.label_encoder.classes_)
logger.info(f"Loaded {self.num_labels} labels")
# 创建线程池
self.executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)
logger.info(f"ThreadPoolExecutor created with {MAX_WORKERS} workers")
# 注册关闭钩子
atexit.register(self._shutdown)
# 模型预热
self._warmup()
self._initialized = True
elapsed_time = time.time() - start_time
logger.info(f"Predictor initialized successfully in {elapsed_time:.2f}s")
except Exception as e:
logger.error(f"Failed to initialize Predictor: {e}")
raise
def _shutdown(self):
"""清理资源"""
logger.info("Shutting down executor...")
self.executor.shutdown(wait=True)
logger.info("Executor shutdown complete")
def _warmup(self):
"""模型预热，确保CUDA初始化"""
logger.info("Warming up model...")
try:
test_text = "测试文本"
self.predict_single(test_text)
logger.info("Model warmup completed")
except Exception as e:
logger.warning(f"Model warmup failed: {e}")
@staticmethod
@lru_cache(maxsize=CACHE_SIZE)
def clean_text(text: str) -> str:
"""文本清洗函数，带缓存"""
if not isinstance(text, str):
return ""
# 只保留中文字符
cleaned_text = re.sub(r'[^\u4e00-\u9fa5]', '', text)
return cleaned_text.strip()
@lru_cache(maxsize=TOKEN_CACHE_SIZE)
def tokenize_text(self, text: str):
"""缓存tokenization结果"""
if not text:
return None
return self.tokenizer(
text,
return_tensors="pt",
truncation=True,
padding=True,
max_length=MAX_LENGTH
)
def predict_single(self, text: str) -> Dict:
"""单个文本预测"""
if not text or not isinstance(text, str):
return {"type": "", "tax": "", "confidence": 0.0, "error": "Invalid input"}
try:
# 清洗文本
cleaned_text = self.clean_text(text)
if not cleaned_text:
logger.warning(f"Empty text after cleaning: {text[:50]}...")
return {"type": "", "tax": "", "confidence": 0.0}
# Tokenization（使用缓存）
inputs = self.tokenize_text(cleaned_text)
if inputs is None:
return {"type": "", "tax": "", "confidence": 0.0}
# 移动到设备
inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
# 推理
with torch.no_grad():
outputs = self.model(**inputs)
probs = torch.softmax(outputs.logits, dim=1).cpu()
top_prob, top_idx = torch.topk(probs, k=1)
# 获取标签
label_idx = top_idx.item()
tax_label = self.label_encoder.inverse_transform([label_idx])[0]
confidence = top_prob.item()
# 尝试解析标签格式
tax_arr = tax_label.split("_")
if len(tax_arr) >= 2:
result = {
"type": tax_arr[0],
"tax": tax_arr[1],
"confidence": round(confidence, 4)
}
else:
result = {
"type": tax_label,
"tax": "",
"confidence": round(confidence, 4)
}
return result
except torch.cuda.OutOfMemoryError as e:
logger.error(f"CUDA OOM error: {e}")
torch.cuda.empty_cache()
return {"type": "", "tax": "", "confidence": 0.0, "error": "GPU memory exhausted"}
except Exception as e:
logger.error(f"Prediction error for text '{text[:50]}...': {e}")
return {"type": "", "tax": "", "confidence": 0.0, "error": str(e)}
def batch_predict(self, texts: List[str], batch_size: int = BATCH_SIZE) -> List[Dict]:
"""批量预测（并发处理）"""
if not texts:
return []
logger.info(f"Processing {len(texts)} texts with batch_size={batch_size}")
start_time = time.time()
results = []
total_batches = (len(texts) + batch_size - 1) // batch_size
for batch_idx in range(total_batches):
batch_start = batch_idx * batch_size
batch_end = min(batch_start + batch_size, len(texts))
batch_texts = texts[batch_start:batch_end]
# 提交批次任务
futures = [
self.executor.submit(self.predict_single, text)
for text in batch_texts
]
# 收集结果
batch_results = [future.result() for future in futures]
results.extend(batch_results)
if batch_idx % 10 == 0 and batch_idx > 0:
logger.info(f"Processed {batch_end}/{len(texts)} texts")
elapsed_time = time.time() - start_time
logger.info(f"Batch prediction completed in {elapsed_time:.2f}s for {len(texts)} texts")
return results
def get_stats(self) -> Dict:
"""获取预测器统计信息"""
return {
"service_name": SERVICE_NAME,
"model_id": MODEL_ID,
"model_dir": MODEL_DIR,
"device": DEVICE,
"num_labels": self.num_labels,
"max_length": MAX_LENGTH,
"cache_size": {
"text_clean": CACHE_SIZE,
"tokenization": TOKEN_CACHE_SIZE
},
"thread_pool": {
"max_workers": MAX_WORKERS,
"batch_size": BATCH_SIZE
}
}
# 初始化预测器（单例）
try:
predictor = Predictor()
logger.info("Predictor instance created successfully")
except Exception as e:
logger.error(f"Failed to create predictor instance: {e}")
raise
@app.route('/predicts', methods=['POST'])
def predicts():
"""预测接口"""
try:
# 解析请求
data = request.get_json()
if not data:
return jsonify({
"status": "error",
"error": "Invalid request, JSON body required"
}), 400
if 'product_names' not in data:
return jsonify({
"status": "error",
"error": "Invalid request, 'product_names' is required"
}), 400
product_names = data['product_names']
if not isinstance(product_names, list):
product_names = [product_names]
# 限制最大处理数量
max_items = data.get('max_items', 1000)
if len(product_names) > max_items:
logger.warning(f"Too many items: {len(product_names)} > {max_items}, truncating")
product_names = product_names[:max_items]
# 去重（可选）
if data.get('deduplicate', False):
product_names = list(dict.fromkeys(product_names))
logger.info(f"After deduplication: {len(product_names)} unique texts")
# 批量预测
results = predictor.batch_predict(product_names)
# 构建响应
response = {
"status": "success",
"predictions": results,
"metadata": {
"total": len(results),
"timestamp": datetime.now().isoformat(),
"service": SERVICE_NAME,
"device": DEVICE
}
}
return jsonify(response)
except Exception as e:
logger.error(f"Prediction endpoint error: {e}", exc_info=True)
return jsonify({
"status": "error",
"error": f"Internal server error: {str(e)}"
}), 500
@app.route('/predict', methods=['POST'])
def predict():
"""单条预测接口（简化版）"""
try:
data = request.get_json()
if not data or 'product_name' not in data:
return jsonify({
"status": "error",
"error": "Invalid request, 'product_name' is required"
}), 400
product_name = data['product_name']
result = predictor.predict_single(product_name)
return jsonify({
"status": "success",
"prediction": result
})
except Exception as e:
logger.error(f"Single prediction endpoint error: {e}", exc_info=True)
return jsonify({
"status": "error",
"error": str(e)
}), 500
@app.route('/health', methods=['GET'])
def health_check():
"""健康检查接口"""
try:
stats = predictor.get_stats()
return jsonify({
"status": "healthy",
"timestamp": datetime.now().isoformat(),
"stats": stats
})
except Exception as e:
return jsonify({
"status": "unhealthy",
"error": str(e)
}), 500
@app.route('/stats', methods=['GET'])
def get_stats():
"""获取预测器统计信息"""
try:
stats = predictor.get_stats()
return jsonify({
"status": "success",
"stats": stats
})
except Exception as e:
return jsonify({
"status": "error",
"error": str(e)
}), 500
@app.errorhandler(404)
def not_found(error):
return jsonify({
"status": "error",
"error": "Endpoint not found"
}), 404
@app.errorhandler(500)
def internal_error(error):
return jsonify({
"status": "error",
"error": "Internal server error"
}), 500
if __name__ == '__main__':
logger.info(f"Starting Flask application: {SERVICE_NAME}")
logger.info(f"Model path: {MODEL_DIR}")
logger.info(f"Device: {DEVICE}")
logger.info(f"Port: {SERVICE_PORT}")
logger.info(f"Batch size: {BATCH_SIZE}")
logger.info(f"Max workers: {MAX_WORKERS}")
app.run(
host='0.0.0.0',
port=SERVICE_PORT,
threaded=True,
debug=False
)
else:
# 兼容 WSGI 标准（如 Gunicorn）
application = app
logger.info(f"Application loaded for WSGI server: {SERVICE_NAME}")

825
services/tax/train.py Normal file
View File

@ -0,0 +1,825 @@
"""
RoBERTa-wwm-ext-large 商品分类微调脚本 (ModelScope版)
硬件配置: 8核CPU / 32GB内存 / 24GB显存
模型来源: ModelScope
"""
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.utils.class_weight import compute_class_weight
from sklearn.metrics import accuracy_score, f1_score, classification_report
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
Trainer,
TrainingArguments,
EarlyStoppingCallback,
set_seed,
TrainerCallback
)
from transformers.modeling_outputs import SequenceClassifierOutput
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from tqdm import tqdm
import warnings
import re
import joblib
import os
import time
import glob
from datetime import datetime
from torch.utils.tensorboard import SummaryWriter
from safetensors.torch import save_file as safe_save_file
# ==================== ModelScope 导入 ====================
try:
from modelscope.hub.snapshot_download import snapshot_download
MODELSCOPE_AVAILABLE = True
print("✅ ModelScope 已安装")
except ImportError:
MODELSCOPE_AVAILABLE = False
print("⚠️ ModelScope 未安装，将使用HuggingFace缓存")
print(" 安装命令: pip install modelscope")
# ==================== 1. 参数配置 ====================
class Config:
# ===== 模型配置 =====
MODEL_ID = "dienstag/chinese-roberta-wwm-ext"
# 本地缓存目录
CACHE_DIR = "./modelscope_cache"
# 是否使用ModelScope下载
USE_MODELSCOPE = True
# ===== 训练参数 =====
MAX_LENGTH = 64
BATCH_SIZE = 32
GRADIENT_ACCUMULATION_STEPS = 2
NUM_EPOCHS = 5
LEARNING_RATE = 2e-5
WARMUP_RATIO = 0.1
WEIGHT_DECAY = 0.01
FP16 = True
EARLY_STOPPING_PATIENCE = 3
SAVE_TOTAL_LIMIT = 5
# 日志和保存步数
LOGGING_STEPS = 50
EVAL_SAVE_STEPS = 200
# 数据加载
NUM_WORKERS = 6
TEST_SIZE = 0.2
# 路径
DATA_FILE = "train_v4.xlsx"
OUTPUT_DIR = "./results/roberta_large"
LOG_DIR = "./logs/roberta_large"
SAVE_DIR = "./saved_model/roberta_large"
# 随机种子
SEED = 42
# 设备
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# 评估指标配置
METRIC_FOR_BEST_MODEL = "f1_macro"
GREATER_IS_BETTER = True
# ==================== 2. ModelScope 模型加载 ====================
def download_model_from_modelscope(model_id, cache_dir):
"""从ModelScope下载模型"""
if not MODELSCOPE_AVAILABLE:
raise ImportError("请先安装 modelscope: pip install modelscope")
print(f"📥 从ModelScope下载模型: {model_id}")
print(f"📁 缓存目录: {cache_dir}")
try:
model_dir = snapshot_download(
model_id,
cache_dir=cache_dir,
revision='master'
)
print(f"✅ 模型下载完成: {model_dir}")
return model_dir
except Exception as e:
print(f"⚠️ ModelScope下载失败: {e}")
print(" 尝试使用HuggingFace缓存...")
return None
def get_model_path(model_id, cache_dir, use_modelscope=True):
"""获取模型路径"""
if use_modelscope and MODELSCOPE_AVAILABLE:
try:
local_dir = download_model_from_modelscope(model_id, cache_dir)
if local_dir and os.path.exists(local_dir):
return local_dir
except Exception as e:
print(f"⚠️ ModelScope加载失败: {e}")
print(f"🔄 使用HuggingFace加载: {model_id}")
return model_id
# ==================== 3. 工具函数 ====================
def print_gpu_memory():
"""打印GPU显存使用情况"""
if torch.cuda.is_available():
allocated = torch.cuda.memory_allocated() / 1024 ** 3
reserved = torch.cuda.memory_reserved() / 1024 ** 3
max_allocated = torch.cuda.max_memory_allocated() / 1024 ** 3
print(f"📊 GPU显存: 已分配 {allocated:.2f}GB | 已预留 {reserved:.2f}GB | 峰值 {max_allocated:.2f}GB")
def format_time(seconds):
"""格式化时间显示"""
if seconds < 60:
return f"{seconds:.0f}秒"
elif seconds < 3600:
return f"{seconds / 60:.1f}分钟"
else:
return f"{seconds / 3600:.2f}小时"
# ==================== 4. 数据加载 ====================
def load_data(file_path):
"""加载数据文件"""
try:
if file_path.endswith('.csv'):
df = pd.read_csv(file_path)
elif file_path.endswith(('.xlsx', '.xls')):
df = pd.read_excel(file_path)
else:
raise ValueError("不支持的文件格式，请使用 .csv 或 .xlsx")
print(f"📄 原始数据列: {df.columns.tolist()}")
# 自动识别文本列
if 'sentence' not in df.columns:
possible_text_cols = ['sentence', 'content', 'name', 'product_name',
'goods_name', 'title', 'Title', '商品名称', '商品全称']
for col in possible_text_cols:
if col in df.columns:
df.rename(columns={col: 'sentence'}, inplace=True)
break
# 自动识别标签列
if 'label' not in df.columns:
possible_label_cols = ['label', 'class', 'type', 'cate',
'Type', '分类', '类别']
for col in possible_label_cols:
if col in df.columns:
df.rename(columns={col: 'label'}, inplace=True)
break
assert 'sentence' in df.columns and 'label' in df.columns, \
"数据必须包含'sentence'(文本)和'label'(标签)列"
df = df.dropna(subset=['sentence', 'label']).reset_index(drop=True)
print(f"✅ 数据加载成功 | 样本量: {len(df)} | 分类数: {df['label'].nunique()}")
return df
except Exception as e:
warnings.warn(f"❌ 数据加载失败: {str(e)}")
raise
# ==================== 5. 数据清洗 ====================
def clean_chinese_text(text):
"""优化版清洗：保留中文字符、英文字母、数字、空格和重要符号"""
if not isinstance(text, str):
return ""
cleaned = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s\-/\.\(\)\（\）\·]', ' ', text)
cleaned = re.sub(r'\s+', ' ', cleaned).strip()
return cleaned
def clean_batch_texts(texts, show_sample=True):
"""批量清洗文本"""
print("🧼 开始清洗文本数据...")
cleaned = [clean_chinese_text(t) for t in texts]
original_len = sum(len(t) for t in texts)
cleaned_len = sum(len(t) for t in cleaned)
if original_len > 0:
print(f" 原始总字符: {original_len:,} | 清洗后总字符: {cleaned_len:,} | 缩减: {(1 - cleaned_len / original_len) * 100:.1f}%")
if show_sample and len(texts) > 0:
print("\n📝 清洗样例:")
for i in range(min(3, len(texts))):
print(f" 原始: {texts[i][:50]}{'...' if len(texts[i]) > 50 else ''}")
print(f" 清洗: {cleaned[i][:50]}{'...' if len(cleaned[i]) > 50 else ''}")
print()
return cleaned
# ==================== 6. 数据集类 ====================
class TextDataset(Dataset):
"""动态编码数据集"""
def __init__(self, texts, labels, tokenizer, max_length):
self.texts = texts
self.labels = labels
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
encoding = self.tokenizer(
self.texts[idx],
max_length=self.max_length,
padding=False,
truncation=True,
return_tensors=None
)
return {
"input_ids": encoding["input_ids"],
"attention_mask": encoding["attention_mask"],
"labels": self.labels[idx]
}
# ==================== 7. Collate函数 ====================
def collate_fn(batch, tokenizer):
"""动态填充批处理"""
input_ids = [item["input_ids"] for item in batch]
attention_mask = [item["attention_mask"] for item in batch]
labels = [item["labels"] for item in batch]
padded = tokenizer.pad(
{"input_ids": input_ids, "attention_mask": attention_mask},
padding=True,
return_tensors="pt"
)
return {
"input_ids": padded["input_ids"],
"attention_mask": padded["attention_mask"],
"labels": torch.tensor(labels, dtype=torch.long)
}
# ==================== 8. 评估指标 ====================
def compute_metrics(eval_pred):
"""计算多种评估指标"""
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
return {
"accuracy": accuracy_score(labels, predictions),
"f1_macro": f1_score(labels, predictions, average="macro"),
"f1_weighted": f1_score(labels, predictions, average="weighted"),
"f1_micro": f1_score(labels, predictions, average="micro"),
}
# ==================== 9. 类别权重 ====================
def compute_class_weights(labels, num_classes):
"""计算类别权重"""
classes = np.unique(labels)
weights = compute_class_weight("balanced", classes=classes, y=labels)
weight_tensor = torch.ones(num_classes, dtype=torch.float)
for cls, weight in zip(classes, weights):
weight_tensor[cls] = weight
return weight_tensor
# ==================== 10. 带权重损失的自定义模型 ====================
class CustomClassificationModel(torch.nn.Module):
"""包装原始模型，支持自定义损失函数"""
def __init__(self, base_model, class_weights=None):
super().__init__()
self.base_model = base_model
self.num_labels = base_model.config.num_labels
self.config = base_model.config
if class_weights is not None:
self.loss_fct = torch.nn.CrossEntropyLoss(weight=class_weights)
else:
self.loss_fct = torch.nn.CrossEntropyLoss()
def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
outputs = self.base_model(
input_ids=input_ids,
attention_mask=attention_mask,
**kwargs
)
logits = outputs.logits
loss = None
if labels is not None:
loss = self.loss_fct(logits, labels)
return SequenceClassifierOutput(
loss=loss,
logits=logits,
hidden_states=outputs.hidden_states if hasattr(outputs, 'hidden_states') else None,
attentions=outputs.attentions if hasattr(outputs, 'attentions') else None
)
# ==================== 11. 自定义保存回调（修复非连续张量问题） ====================
class CustomSaveCallback(TrainerCallback):
"""
自定义保存回调：在保存前修复非连续张量
"""
def on_save(self, args, state, control, model, tokenizer=None, **kwargs):
if state.is_world_process_zero:
# 获取当前保存路径
checkpoint_dir = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
# 如果是包装模型，获取基础模型
if hasattr(model, 'base_model'):
model_to_save = model.base_model
else:
model_to_save = model
# 确保所有参数连续
for param in model_to_save.parameters():
if not param.is_contiguous():
param.data = param.data.contiguous()
# 保存模型（使用pickle格式）
try:
model_to_save.save_pretrained(checkpoint_dir)
if tokenizer:
tokenizer.save_pretrained(checkpoint_dir)
print(f"✅ 检查点已保存: {checkpoint_dir}")
except Exception as e:
print(f"⚠️ 保存检查点失败: {e}")
# 备用方案：手动保存
self._manual_save(model_to_save, tokenizer, checkpoint_dir)
def _manual_save(self, model, tokenizer, save_dir):
"""手动保存模型"""
try:
os.makedirs(save_dir, exist_ok=True)
model.config.save_pretrained(save_dir)
state_dict = model.state_dict()
contiguous_state_dict = {}
for key, tensor in state_dict.items():
if not tensor.is_contiguous():
contiguous_state_dict[key] = tensor.contiguous()
else:
contiguous_state_dict[key] = tensor
torch.save(contiguous_state_dict, os.path.join(save_dir, "pytorch_model.bin"))
if tokenizer:
tokenizer.save_pretrained(save_dir)
print(f"✅ 检查点已保存 (手动): {save_dir}")
except Exception as e:
print(f"❌ 保存失败: {e}")
# ==================== 12. TensorBoard回调 ====================
class TensorBoardCallback(TrainerCallback):
"""TensorBoard回调"""
def __init__(self, log_dir):
self.writer = SummaryWriter(log_dir)
def on_log(self, args, state, control, logs=None, **kwargs):
if logs:
step = state.global_step
for key, value in logs.items():
if isinstance(value, (int, float)):
self.writer.add_scalar(key, value, step)
def on_train_end(self, args, state, control, **kwargs):
self.writer.close()
# ==================== 13. 模型初始化 ====================
def init_model(num_labels, class_weights=None):
"""初始化模型和分词器"""
model_path = get_model_path(
Config.MODEL_ID,
Config.CACHE_DIR,
Config.USE_MODELSCOPE
)
print(f"\n🔄 加载模型: {model_path}")
start_time = time.time()
tokenizer = AutoTokenizer.from_pretrained(
model_path,
trust_remote_code=True
)
base_model = AutoModelForSequenceClassification.from_pretrained(
model_path,
num_labels=num_labels,
trust_remote_code=True
)
# 确保模型参数是连续的（修复非连续张量问题）
base_model = base_model.to(Config.DEVICE)
# 遍历所有参数，确保连续
for param in base_model.parameters():
if not param.is_contiguous():
param.data = param.data.contiguous()
if class_weights is not None:
class_weights = class_weights.to(Config.DEVICE)
model = CustomClassificationModel(base_model, class_weights)
else:
model = base_model
load_time = time.time() - start_time
print(f"✅ 模型加载完成，耗时: {format_time(load_time)}")
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"📊 模型参数总量: {total_params:,} ({total_params / 1e6:.1f}M)")
print(f"📊 可训练参数: {trainable_params:,} ({trainable_params / 1e6:.1f}M)")
return tokenizer, model
# ==================== 14. 训练配置 ====================
def get_training_args():
"""获取训练参数配置"""
return TrainingArguments(
output_dir=Config.OUTPUT_DIR,
num_train_epochs=Config.NUM_EPOCHS,
per_device_train_batch_size=Config.BATCH_SIZE,
per_device_eval_batch_size=Config.BATCH_SIZE * 2,
gradient_accumulation_steps=Config.GRADIENT_ACCUMULATION_STEPS,
learning_rate=Config.LEARNING_RATE,
warmup_ratio=Config.WARMUP_RATIO,
weight_decay=Config.WEIGHT_DECAY,
# 日志和保存
logging_dir=Config.LOG_DIR,
logging_steps=Config.LOGGING_STEPS,
logging_first_step=True,
eval_strategy="steps",
eval_steps=Config.EVAL_SAVE_STEPS,
save_strategy="steps",
save_steps=Config.EVAL_SAVE_STEPS,
save_total_limit=Config.SAVE_TOTAL_LIMIT,
load_best_model_at_end=True,
metric_for_best_model=Config.METRIC_FOR_BEST_MODEL,
greater_is_better=Config.GREATER_IS_BETTER,
# 性能优化
fp16=Config.FP16,
dataloader_num_workers=Config.NUM_WORKERS,
dataloader_pin_memory=True,
# 其他
report_to="tensorboard",
seed=Config.SEED,
push_to_hub=False,
lr_scheduler_type="linear",
remove_unused_columns=False,
dataloader_drop_last=False,
# ✅ 关键修复：禁用 safetensors
save_safetensors=False,
)
# ==================== 15. 保存模型 ====================
def save_model(model, tokenizer, save_dir):
"""保存模型"""
print(f"\n💾 保存模型到: {save_dir}")
os.makedirs(save_dir, exist_ok=True)
if hasattr(model, 'base_model'):
model_to_save = model.base_model
else:
model_to_save = model
# 确保所有参数连续
for param in model_to_save.parameters():
if not param.is_contiguous():
param.data = param.data.contiguous()
try:
model_to_save.save_pretrained(save_dir)
tokenizer.save_pretrained(save_dir)
print(f"✅ 模型已保存到: {save_dir}")
return True
except Exception as e:
print(f"⚠️ 保存失败: {e}")
print(" 尝试使用手动方式保存...")
# 备用方案：手动保存
try:
model_to_save.config.save_pretrained(save_dir)
state_dict = model_to_save.state_dict()
contiguous_state_dict = {}
for key, tensor in state_dict.items():
if not tensor.is_contiguous():
contiguous_state_dict[key] = tensor.contiguous()
else:
contiguous_state_dict[key] = tensor
torch.save(contiguous_state_dict, os.path.join(save_dir, "pytorch_model.bin"))
tokenizer.save_pretrained(save_dir)
print(f"✅ 模型已保存到: {save_dir}")
return True
except Exception as e:
print(f"❌ 保存失败: {e}")
return False
# ==================== 16. 预测函数 ====================
@torch.no_grad()
def batch_predict(texts, model, tokenizer, label_map, top_k=3, batch_size=64):
"""批量预测"""
model.eval()
all_results = []
cleaned_texts = [clean_chinese_text(t) for t in texts]
for i in tqdm(range(0, len(cleaned_texts), batch_size), desc="预测中"):
batch = cleaned_texts[i:i + batch_size]
inputs = tokenizer(
batch,
return_tensors="pt",
truncation=True,
padding=True,
max_length=Config.MAX_LENGTH
).to(Config.DEVICE)
outputs = model(**inputs)
if hasattr(outputs, 'logits'):
logits = outputs.logits
else:
logits = outputs
probs = torch.softmax(logits, dim=1).cpu()
for prob in probs:
top_probs, top_indices = torch.topk(prob, k=min(top_k, len(label_map)))
results = [
{
"category": label_map[idx.item()],
"confidence": prob.item()
}
for prob, idx in zip(top_probs, top_indices)
]
all_results.append(results)
return all_results
# ==================== 17. 模型验证 ====================
def validate_model(model, tokenizer, sample_texts):
"""快速验证模型"""
print("\n🔍 模型验证...")
try:
inputs = tokenizer(
sample_texts,
return_tensors="pt",
padding=True,
truncation=True,
max_length=Config.MAX_LENGTH
).to(Config.DEVICE)
outputs = model(**inputs)
if hasattr(outputs, 'logits'):
logits = outputs.logits
else:
logits = outputs
print(f"✅ 模型验证通过，输出形状: {logits.shape}")
return True
except Exception as e:
print(f"❌ 模型验证失败: {e}")
return False
# ==================== 18. 数据统计 ====================
def print_data_stats(df):
"""打印数据统计信息"""
print("\n" + "=" * 50)
print("📊 数据统计")
print("=" * 50)
print(f"总样本量: {len(df)}")
print(f"分类数: {df['label'].nunique()}")
label_counts = df['label'].value_counts()
print(f"\n📈 类别分布 (Top 10):")
for label, count in label_counts.head(10).items():
pct = count / len(df) * 100
bar = "█" * int(pct / 2)
print(f" {label}: {count} ({pct:.1f}%) {bar}")
tail_threshold = 10
tail_count = (label_counts <= tail_threshold).sum()
tail_samples = label_counts[label_counts <= tail_threshold].sum()
print(f"\n📉 长尾类别 (≤{tail_threshold}条): {tail_count}个类别, 共{tail_samples}条样本")
text_lengths = df['sentence'].str.len()
print(f"\n📏 文本长度统计:")
print(f" 平均: {text_lengths.mean():.1f}")
print(f" 中位数: {text_lengths.median():.0f}")
print(f" 最大: {text_lengths.max()}")
print(f" 最小: {text_lengths.min()}")
print("=" * 50 + "\n")
# ==================== 19. 主函数 ====================
def main():
"""主训练流程"""
print("\n" + "=" * 60)
print("🚀 RoBERTa-wwm-ext-large 商品分类微调 (ModelScope版)")
print(f"📅 启动时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
print(f"💻 设备: {Config.DEVICE}")
print(f"📦 模型来源: {'ModelScope' if Config.USE_MODELSCOPE else 'HuggingFace'}")
if torch.cuda.is_available():
print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
print(f"💾 显存: {torch.cuda.get_device_properties(0).total_memory / 1024 ** 3:.1f}GB")
print("=" * 60)
total_start = time.time()
set_seed(Config.SEED)
os.makedirs(Config.OUTPUT_DIR, exist_ok=True)
os.makedirs(Config.LOG_DIR, exist_ok=True)
os.makedirs(Config.SAVE_DIR, exist_ok=True)
os.makedirs(Config.CACHE_DIR, exist_ok=True)
# ===== 1. 加载数据 =====
print(f"\n📂 加载数据: {Config.DATA_FILE}")
df = load_data(Config.DATA_FILE)
# ===== 2. 数据清洗 =====
df['sentence_cleaned'] = clean_batch_texts(df['sentence'].tolist())
df = df[df['sentence_cleaned'].str.len() > 0].reset_index(drop=True)
print(f"✅ 清洗后剩余: {len(df)} 条")
# ===== 3. 数据统计 =====
print_data_stats(df)
# ===== 4. 标签编码 =====
print("🏷️ 处理标签...")
label_encoder = LabelEncoder()
df['label_id'] = label_encoder.fit_transform(df['label'])
label_map = {i: label for i, label in enumerate(label_encoder.classes_)}
num_labels = len(label_map)
print(f"✅ 类别数: {num_labels}")
print(f" 标签示例: {list(label_map.items())[:5]}")
joblib.dump(label_encoder, "label_encoder_roberta_large.pkl")
print("✅ 标签编码器已保存: label_encoder_roberta_large.pkl")
# ===== 5. 类别权重 =====
print("\n⚖️ 计算类别权重...")
class_weights = compute_class_weights(df['label_id'].values, num_labels)
print(f" 权重范围: {class_weights.min():.3f} - {class_weights.max():.3f}")
print(f" 平均权重: {class_weights.mean():.3f}")
# ===== 6. 划分数据集 =====
print("\n📊 划分数据集...")
train_df, eval_df = train_test_split(
df,
test_size=Config.TEST_SIZE,
random_state=Config.SEED,
stratify=df["label_id"]
)
print(f" 训练集: {len(train_df)} 条")
print(f" 验证集: {len(eval_df)} 条")
# ===== 7. 初始化模型 =====
tokenizer, model = init_model(num_labels, class_weights)
model.to(Config.DEVICE)
print(f"✅ 模型已加载到: {Config.DEVICE}")
print_gpu_memory()
# ===== 8. 模型验证 =====
sample_texts = ["九阳电饭煲家用", "苹果手机iPhone", "耐克运动鞋"]
if not validate_model(model, tokenizer, sample_texts):
print("⚠️ 模型验证失败，继续训练...")
# ===== 9. 准备数据集 =====
print("\n📦 准备数据集...")
train_dataset = TextDataset(
train_df['sentence_cleaned'].tolist(),
train_df['label_id'].tolist(),
tokenizer,
Config.MAX_LENGTH
)
eval_dataset = TextDataset(
eval_df['sentence_cleaned'].tolist(),
eval_df['label_id'].tolist(),
tokenizer,
Config.MAX_LENGTH
)
# ===== 10. 配置训练器 =====
print("\n⚙️ 配置训练器...")
training_args = get_training_args()
def data_collator(batch):
return collate_fn(batch, tokenizer)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
data_collator=data_collator,
compute_metrics=compute_metrics,
callbacks=[
EarlyStoppingCallback(early_stopping_patience=Config.EARLY_STOPPING_PATIENCE),
TensorBoardCallback(Config.LOG_DIR),
CustomSaveCallback(), # ✅ 添加自定义保存回调
]
)
# ===== 11. 断点续训 =====
checkpoint_dir = None
checkpoints = glob.glob(os.path.join(Config.OUTPUT_DIR, "checkpoint-*"))
if checkpoints:
latest_checkpoint = max(checkpoints, key=os.path.getctime)
print(f"📂 发现检查点: {latest_checkpoint}")
checkpoint_dir = latest_checkpoint
# ===== 12. 开始训练 =====
print("\n" + "=" * 60)
print("🚀 开始训练")
print("=" * 60)
train_start = time.time()
if checkpoint_dir:
trainer.train(resume_from_checkpoint=checkpoint_dir)
else:
trainer.train()
train_time = time.time() - train_start
print(f"\n✅ 训练完成，耗时: {format_time(train_time)}")
# ===== 13. 保存模型 =====
save_model(model, tokenizer, Config.SAVE_DIR)
# ===== 14. 最终评估 =====
print("\n📊 最终评估...")
eval_results = trainer.evaluate()
print(f" Eval Loss: {eval_results['eval_loss']:.4f}")
print(f" Eval Accuracy: {eval_results.get('eval_accuracy', 0):.4f}")
print(f" Eval F1 Macro: {eval_results.get('eval_f1_macro', 0):.4f}")
print(f" Eval F1 Weighted: {eval_results.get('eval_f1_weighted', 0):.4f}")
# ===== 15. 推理测试 =====
print("\n🧪 推理测试...")
test_samples = [
"九阳3升家用多功能电饭煲F30S-S160",
"苹果iPhone 15 Pro Max 512GB 黑色钛金属",
"海尔双开门冰箱BCD-600WGHSS19B8U1",
"Lining李宁女款长裤 AKYV092-1 黑色",
"珀莱雅双抗精华美白特证版补水保湿套装50ml礼盒",
]
predictions = batch_predict(test_samples, model, tokenizer, label_map, top_k=3)
print("\n📋 预测结果:")
print("-" * 80)
for sample, preds in zip(test_samples, predictions):
print(f"输入: {sample}")
for i, p in enumerate(preds):
print(f" Top{i + 1}: {p['category']} (置信度: {p['confidence']:.4f})")
print("-" * 80)
# ===== 16. 总结 =====
total_time = time.time() - total_start
print("\n" + "=" * 60)
print("🎉 训练完成!")
print(f"📅 结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
print(f"⏱️ 总耗时: {format_time(total_time)}")
print(f"⏱️ 训练耗时: {format_time(train_time)}")
print(f"📊 最佳模型指标: {Config.METRIC_FOR_BEST_MODEL} = {trainer.state.best_metric:.4f}")
print("=" * 60)
# ==================== 20. 入口 ====================
if __name__ == "__main__":
main()