diff --git a/.env b/.env new file mode 100644 index 0000000..d058eee --- /dev/null +++ b/.env @@ -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 diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..9fca517 --- /dev/null +++ b/Dockerfile @@ -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"] \ No newline at end of file diff --git a/base/Dockerfile b/base/Dockerfile new file mode 100644 index 0000000..d17d40f --- /dev/null +++ b/base/Dockerfile @@ -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 + + diff --git a/base/requirements.txt b/base/requirements.txt new file mode 100644 index 0000000..a3ae07a --- /dev/null +++ b/base/requirements.txt @@ -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 \ No newline at end of file diff --git a/deploy.sh b/deploy.sh new file mode 100644 index 0000000..06af75d --- /dev/null +++ b/deploy.sh @@ -0,0 +1,2 @@ +# 3. 首次启动（会下载模型到本地） +docker-compose up -d diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..f86bcf6 --- /dev/null +++ b/docker-compose.yml @@ -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] \ No newline at end of file diff --git a/scripts/download_models.py b/scripts/download_models.py new file mode 100644 index 0000000..700e2ab --- /dev/null +++ b/scripts/download_models.py @@ -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() \ No newline at end of file diff --git a/scripts/generate_supervisor_conf.py b/scripts/generate_supervisor_conf.py new file mode 100644 index 0000000..dd6667e --- /dev/null +++ b/scripts/generate_supervisor_conf.py @@ -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() \ No newline at end of file diff --git a/services/address/app.py b/services/address/app.py new file mode 100644 index 0000000..9ce48e9 --- /dev/null +++ b/services/address/app.py @@ -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}") \ No newline at end of file diff --git a/services/address/train.py b/services/address/train.py new file mode 100644 index 0000000..8b8d566 --- /dev/null +++ b/services/address/train.py @@ -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") \ No newline at end of file diff --git a/services/tax/app.py b/services/tax/app.py new file mode 100644 index 0000000..6123a45 --- /dev/null +++ b/services/tax/app.py @@ -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}") \ No newline at end of file diff --git a/services/tax/train.py b/services/tax/train.py new file mode 100644 index 0000000..31e5248 --- /dev/null +++ b/services/tax/train.py @@ -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() \ No newline at end of file