question_classifier.py

#!/usr/bin/env python3
"""
问题分类脚本
使用LLM为GSM8K问题进行分类，以便进行更细粒度的角色性能分析
"""

import json
import pandas as pd
import os
import logging
from typing import Dict, List, Any
from dataclasses import dataclass
from datetime import datetime

from src.models import get_model, InferenceConfig
from src.datasets import GSM8KProcessor

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 预定义的数学问题类别
MATH_CATEGORIES = [
    "arithmetic_operations",      # 基础算术运算
    "word_problems_money",        # 金钱相关应用题
    "word_problems_time",         # 时间相关应用题  
    "word_problems_distance",     # 距离/速度相关应用题
    "percentage_problems",        # 百分比问题
    "ratio_proportion",           # 比例和比率问题
    "algebraic_thinking",         # 代数思维问题
    "multi_step_reasoning",       # 多步推理问题
    "pattern_sequence",           # 模式和序列问题
    "geometry_measurement",       # 几何和测量问题
    "probability_statistics",     # 概率统计问题
    "logical_reasoning",          # 逻辑推理问题
]

CLASSIFICATION_PROMPT = """你是一位数学教育专家，请将下面的数学问题分类到最合适的类别中。

可选类别及其定义：
1. arithmetic_operations - 基础算术运算：主要涉及加减乘除等基本运算
2. word_problems_money - 金钱相关应用题：涉及价格、购买、找零、成本等金钱计算
3. word_problems_time - 时间相关应用题：涉及时间计算、工作效率、工期等
4. word_problems_distance - 距离/速度相关应用题：涉及速度、距离、时间的关系
5. percentage_problems - 百分比问题：涉及百分比计算、折扣、增长率等
6. ratio_proportion - 比例和比率问题：涉及比例关系、比率计算
7. algebraic_thinking - 代数思维问题：需要建立方程或使用代数方法
8. multi_step_reasoning - 多步推理问题：需要多个步骤的复杂推理
9. pattern_sequence - 模式和序列问题：涉及数列、规律发现
10. geometry_measurement - 几何和测量问题：涉及形状、面积、体积等
11. probability_statistics - 概率统计问题：涉及概率计算、统计分析
12. logical_reasoning - 逻辑推理问题：主要考查逻辑思维能力

问题：{question}

请仔细分析这个问题的特点，选择最符合的类别。如果问题涉及多个类别，请选择最主要的一个。

请只输出类别名称（英文），不要包含其他内容。

分类结果："""

class QuestionClassifier:
    """问题分类器"""
    
    def __init__(self, model_name: str = "llama3.1-8b"):
        self.model_name = model_name
        self.model = get_model(model_name)
        self.dataset_processor = GSM8KProcessor()
        self.inference_config = InferenceConfig(
            num_samples=1,
            temperature=0.3,  # 较低的温度以获得更一致的分类结果
            top_p=0.8,
            max_new_tokens=50  # 分类结果很短
        )
        
        # 创建输出目录
        self.output_dir = "question_classification_results"
        os.makedirs(self.output_dir, exist_ok=True)
    
    def classify_question(self, question: str, question_id: int) -> Dict[str, Any]:
        """对单个问题进行分类"""
        try:
            # 构建分类prompt
            prompt = CLASSIFICATION_PROMPT.format(question=question)
            
            # 生成分类结果
            response = self.model.generate_response(prompt, self.inference_config)
            
            # 提取分类结果
            category = self.extract_category(response)
            
            return {
                "question_id": question_id,
                "question": question,
                "classification_response": response.strip(),
                "category": category,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            logger.error(f"Error classifying question {question_id}: {e}")
            return {
                "question_id": question_id,
                "question": question,
                "classification_response": "",
                "category": "unknown",
                "timestamp": datetime.now().isoformat(),
                "error": str(e)
            }
    
    def extract_category(self, response: str) -> str:
        """从响应中提取分类类别"""
        response = response.strip().lower()
        
        # 检查响应中是否包含预定义的类别
        for category in MATH_CATEGORIES:
            if category.lower() in response:
                return category
        
        # 如果没有找到匹配的类别，尝试更灵活的匹配
        category_mapping = {
            "arithmetic": "arithmetic_operations",
            "money": "word_problems_money",
            "time": "word_problems_time", 
            "distance": "word_problems_distance",
            "speed": "word_problems_distance",
            "percentage": "percentage_problems",
            "percent": "percentage_problems",
            "ratio": "ratio_proportion",
            "proportion": "ratio_proportion",
            "algebra": "algebraic_thinking",
            "equation": "algebraic_thinking",
            "multi": "multi_step_reasoning",
            "pattern": "pattern_sequence",
            "sequence": "pattern_sequence",
            "geometry": "geometry_measurement",
            "measurement": "geometry_measurement",
            "probability": "probability_statistics",
            "statistics": "probability_statistics",
            "logical": "logical_reasoning",
            "logic": "logical_reasoning"
        }
        
        for keyword, category in category_mapping.items():
            if keyword in response:
                return category
        
        # 如果仍然没有匹配，返回unknown
        logger.warning(f"Could not extract category from response: {response}")
        return "unknown"
    
    def classify_dataset(self, num_samples: int = 200) -> pd.DataFrame:
        """对数据集中的问题进行分类"""
        logger.info(f"开始对GSM8K问题进行分类...")
        logger.info(f"使用模型: {self.model_name}")
        logger.info(f"样本数量: {num_samples}")
        
        # 采样数据
        dataset_list = list(self.dataset_processor.dataset)
        samples = dataset_list[:num_samples] if num_samples < len(dataset_list) else dataset_list
        logger.info(f"实际使用样本数: {len(samples)}")
        
        classification_results = []
        
        # 对每个问题进行分类
        for idx, sample in enumerate(samples):
            if (idx + 1) % 10 == 0:
                logger.info(f"进度: {idx + 1}/{len(samples)}")
            
            result = self.classify_question(sample["question"], idx)
            classification_results.append(result)
        
        # 转换为DataFrame
        df = pd.DataFrame(classification_results)
        
        # 保存分类结果
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_path = os.path.join(self.output_dir, f"question_categories_{timestamp}.csv")
        df.to_csv(csv_path, index=False, encoding='utf-8')
        
        json_path = os.path.join(self.output_dir, f"question_categories_{timestamp}.json")
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(classification_results, f, indent=2, ensure_ascii=False)
        
        logger.info(f"分类结果已保存到: {csv_path}")
        logger.info(f"分类结果已保存到: {json_path}")
        
        return df
    
    def generate_classification_summary(self, df: pd.DataFrame) -> Dict[str, Any]:
        """生成分类汇总报告"""
        logger.info("\n生成分类汇总报告...")
        
        # 统计每个类别的数量
        category_counts = df['category'].value_counts().to_dict()
        category_percentages = df['category'].value_counts(normalize=True).to_dict()
        
        # 创建汇总信息
        summary = {
            "timestamp": datetime.now().strftime("%Y%m%d_%H%M%S"),
            "model_name": self.model_name,
            "total_questions": len(df),
            "categories_found": len(category_counts),
            "category_distribution": {
                category: {
                    "count": count,
                    "percentage": category_percentages.get(category, 0) * 100
                }
                for category, count in category_counts.items()
            }
        }
        
        # 保存汇总报告
        summary_json_path = os.path.join(self.output_dir, f"classification_summary_{summary['timestamp']}.json")
        with open(summary_json_path, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)
        
        # 打印汇总信息
        logger.info("\n问题类别分布:")
        for category, info in sorted(summary["category_distribution"].items(), 
                                   key=lambda x: x[1]["count"], reverse=True):
            logger.info(f"{category}: {info['count']} 题 ({info['percentage']:.1f}%)")
        
        logger.info(f"\n分类汇总已保存到: {summary_json_path}")
        
        return summary

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="问题分类")
    parser.add_argument("--model", default="llama3.1-8b", choices=["qwen2.5-3b", "ministral-7b", "gemma-2-2b", "llama3.1-8b"],
                       help="选择使用的模型")
    parser.add_argument("--samples", type=int, default=200, help="使用的样本数量")
    
    args = parser.parse_args()
    
    # 创建分类器
    classifier = QuestionClassifier(model_name=args.model)
    
    # 执行分类
    logger.info("🚀 开始问题分类...")
    df = classifier.classify_dataset(num_samples=args.samples)
    
    # 生成汇总报告
    summary = classifier.generate_classification_summary(df)
    
    logger.info("✅ 问题分类完成!")
    logger.info(f"共分类 {summary['total_questions']} 个问题到 {summary['categories_found']} 个类别")

if __name__ == "__main__":
    main()