#文本分类实战：基于BERT的企业级情感分析引擎完整开发指南

import pandas as pd
import re
from sklearn.model_selection import train_test_split
from typing import List, Tuple

class DataPreprocessor:
    def __init__(self):
        # 实际项目加载哈工大/百度停用词库
        self.stop_words = {'的', '了', '在', '是', '很', '也', '这', '那', '都'}
    
    def clean_text(self, text: str) -> str:
        # 1. 保留中文、英文、数字
        text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9]', ' ', text)
        # 2. 去除多余空格
        text = re.sub(r'\s+', ' ', text).strip()
        return text
    
    def split_data(self, df: pd.DataFrame, text_col: str='text', label_col: str='label') -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        # 分层采样，保证类别分布一致
        train, temp = train_test_split(df, test_size=0.3, random_state=42, stratify=df[label_col])
        val, test = train_test_split(temp, test_size=0.5, random_state=42, stratify=temp[label_col])
        return train, val, test

# 模拟业务数据（实际从CSV/数据库/API加载）
def load_demo_data():
    data = {
        'text': [
            '这个产品质量很好，值得推荐！',
            '服务态度很差，完全不推荐。',
            '物流很快，包装也不错',
            '质量一般，性价比不高',
            '非常满意，下次还会购买',
            '商品与描述不符，很失望'
        ],
        'label': [1, 0, 1, 0, 1, 0]  # 1=正面，0=负面
    }
    return pd.DataFrame(data)

# 运行示例
df = load_demo_data()
preprocessor = DataPreprocessor()
df['cleaned_text'] = df['text'].apply(preprocessor.clean_text)

import random

# 轻量数据增强（实际可安装pip install synonyms提升效果）
def light_augment(text: str, n_aug: int=2) -> List[str]:
    synonyms_dict = {
        "好": ["优秀", "棒", "不错", "良好"],
        "坏": ["差", "糟糕", "不好", "恶劣"],
        "喜欢": ["喜爱", "欣赏", "钟爱", "青睐"],
        "讨厌": ["厌恶", "反感", "嫌弃", "不满"]
    }
    augmented = []
    words = list(text)
    for _ in range(n_aug):
        new_words = words.copy()
        for i, w in enumerate(new_words):
            if w in synonyms_dict and random.random() < 0.2:
                new_words[i] = random.choice(synonyms_dict[w])
        augmented.append(''.join(new_words))
    return augmented

# 划分数据并应用增强（仅对训练集）
train_df, val_df, test_df = preprocessor.split_data(df)
augmented_train = []
for _, row in train_df.iterrows():
    augmented_train.append({'cleaned_text': row['cleaned_text'], 'label': row['label']})
    for aug in light_augment(row['cleaned_text']):
        augmented_train.append({'cleaned_text': aug, 'label': row['label']})
augmented_train_df = pd.DataFrame(augmented_train)

from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    Trainer, TrainingArguments, EarlyStoppingCallback
)
from datasets import Dataset
import torch
from sklearn.metrics import accuracy_score, f1_score

# 初始化预训练模型和Tokenizer
MODEL_NAME = "bert-base-chinese"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)

# 数据Token化
def tokenize_function(examples):
    return tokenizer(
        examples["cleaned_text"],
        truncation=True,
        padding="max_length",
        max_length=128
    )

# 转换为Hugging Face Dataset格式
train_ds = Dataset.from_pandas(augmented_train_df[['cleaned_text', 'label']])
val_ds = Dataset.from_pandas(val_df[['cleaned_text', 'label']])
tokenized_train = train_ds.map(tokenize_function, batched=True).remove_columns(['cleaned_text']).rename_column('label', 'labels')
tokenized_val = val_ds.map(tokenize_function, batched=True).remove_columns(['cleaned_text']).rename_column('label', 'labels')

# 定义评估指标
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = torch.argmax(torch.tensor(logits), dim=-1).numpy()
    return {
        "accuracy": accuracy_score(labels, predictions),
        "f1": f1_score(labels, predictions, average='weighted')
    }

# 训练配置
training_args = TrainingArguments(
    output_dir="./sentiment_model",
    num_train_epochs=5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    weight_decay=0.01,
    logging_dir="./sentiment_model/logs",
    logging_steps=20,
    evaluation_strategy="steps",
    eval_steps=50,
    save_strategy="steps",
    save_steps=50,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    greater_is_better=True,
    save_total_limit=2,
    seed=42,
    fp16=torch.cuda.is_available(),  # GPU混合精度
    report_to=None  # 可选：改为'tensorboard'/'wandb'
)

# 初始化Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=2)]
)

# **实战提示**：取消下方注释，在GPU环境中运行
# trainer.train()
# trainer.save_model("./sentiment_model/best_model")
# tokenizer.save_pretrained("./sentiment_model/best_model")

from sklearn.metrics import classification_report, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt

# **实战提示**：取消下方注释，在加载最佳模型后运行
# test_ds = Dataset.from_pandas(test_df[['cleaned_text', 'label']])
# tokenized_test = test_ds.map(tokenize_function, batched=True).remove_columns(['cleaned_text']).rename_column('label', 'labels')
# preds = trainer.predict(tokenized_test).predictions.argmax(-1)
# print(classification_report(test_df['label'], preds, target_names=['负面', '正面']))
# sns.heatmap(confusion_matrix(test_df['label'], preds), annot=True, cmap='Blues')
# plt.show()

from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 初始化FastAPI
app = FastAPI(title="企业级情感分析API", version="1.0.0")

# 定义请求和响应模型
class SingleRequest(BaseModel):
    text: str
class BatchRequest(BaseModel):
    texts: list[str]
class SentimentResponse(BaseModel):
    text: str
    sentiment: str
    confidence: float

# **实战提示**：取消下方注释，加载最佳模型
# MODEL_PATH = "./sentiment_model/best_model"
# tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
# model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
# model.eval()
# label_map = {0: "负面", 1: "正面"}

# 单条预测
@app.post("/predict", response_model=SentimentResponse)
async def predict(req: SingleRequest):
    # **实战提示**：替换为真实模型预测
    return {
        "text": req.text,
        "sentiment": "正面" if "好" in req.text or "棒" in req.text else "负面",
        "confidence": 0.85
    }

# 批量预测
@app.post("/batch_predict")
async def batch_predict(req: BatchRequest):
    # **实战提示**：替换为真实批量预测
    return [
        {"text": t, "sentiment": "正面" if "好" in t or "棒" in t else "负面", "confidence": 0.85}
        for t in req.texts
    ]

# 健康检查
@app.get("/")
async def health_check():
    return {"status": "healthy"}