451 lines
14 KiB
Python
451 lines
14 KiB
Python
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}") |