349 lines
11 KiB
Python
349 lines
11 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', '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}") |