This commit is contained in:
parent
0d5060489a
commit
de92a67288
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
|
@ -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]
|
||||
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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}")
|
||||
|
|
@ -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")
|
||||
|
|
@ -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}")
|
||||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue