AIPentest/promptinjection/guardtest.py

import pandas as pd
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import pipeline
import os
import time
import json
from datetime import datetime

def load_checkpoint(checkpoint_file):
    """加载断点信息"""
    if os.path.exists(checkpoint_file):
        try:
            with open(checkpoint_file, 'r', encoding='utf-8') as f:
                checkpoint = json.load(f)
            print(f"发现断点文件，从第 {checkpoint['last_processed'] + 1} 条记录开始继续处理")
            return checkpoint
        except Exception as e:
            print(f"读取断点文件失败: {e}")
            return None
    return None

def save_checkpoint(checkpoint_file, last_processed, total_records):
    """保存断点信息"""
    checkpoint = {
        'last_processed': last_processed,
        'total_records': total_records,
        'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    }
    try:
        with open(checkpoint_file, 'w', encoding='utf-8') as f:
            json.dump(checkpoint, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"保存断点文件失败: {e}")

def save_batch_results(results, output_file, mode='w'):
    """保存批量结果到CSV"""
    try:
        results_df = pd.DataFrame(results)
        if mode == 'w':
            results_df.to_csv(output_file, index=False, encoding='utf-8')
        else:  # append mode
            results_df.to_csv(output_file, mode='a', header=False, index=False, encoding='utf-8')
        return True
    except Exception as e:
        print(f"保存结果失败: {e}")
        return False

def main():
    # 设置文件路径,可使用ProtectAI,deepset,metaguard86M
    local_path = "./models/metaguard86M"
    input_file = "deepset_prompt_injection_data.csv"
    output_file = "./result/deepset_86M_results.csv"
    checkpoint_file = "86M_checkpoint.json"
    batch_size = 10  # 每批处理的记录数
    
    # 检查模型路径是否存在
    if not os.path.exists(local_path):
        print(f"错误: 模型路径 {local_path} 不存在")
        return
    
    # 加载断点信息
    checkpoint = load_checkpoint(checkpoint_file)
    start_idx = 0
    if checkpoint:
        start_idx = checkpoint['last_processed'] + 1
    
    # 加载模型和分词器
    print("正在加载模型和分词器...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(local_path)
        model = AutoModelForSequenceClassification.from_pretrained(local_path)
        
        # 创建分类器
        classifier = pipeline(
            "text-classification",
            model=model,
            tokenizer=tokenizer,
            truncation=True,
            max_length=512,
        )
        print("模型加载成功!")
    except Exception as e:
        print(f"模型加载失败: {e}")
        return
    
    # 读取CSV文件
    if not os.path.exists(input_file):
        print(f"错误: 输入文件 {input_file} 不存在")
        return
    
    try:
        print(f"正在读取 {input_file}...")
        df = pd.read_csv(input_file)
        print(f"成功读取 {len(df)} 行数据")
        
        # 检查必要的列是否存在
        if 'text' not in df.columns:
            print("错误: CSV文件中缺少 'text' 列")
            return
        
        # 如果没有label列，创建一个空的
        if 'label' not in df.columns:
            df['label'] = ''
            
    except Exception as e:
        print(f"读取CSV文件失败: {e}")
        return
    
    # 如果是断点续传，检查总记录数是否一致
    if checkpoint and checkpoint['total_records'] != len(df):
        print(f"警告: 当前文件记录数({len(df)})与断点记录数({checkpoint['total_records']})不一致")
        response = input("是否要重新开始处理? (y/n): ")
        if response.lower() == 'y':
            start_idx = 0
            if os.path.exists(output_file):
                os.remove(output_file)
    
    # 如果从头开始，清空之前的输出文件
    if start_idx == 0 and os.path.exists(output_file):
        os.remove(output_file)
        print("已清空之前的输出文件")
    
    # 创建结果列表
    batch_results = []
    total_processed = start_idx
    
    # 记录分类开始时间
    start_time = time.time()
    start_datetime = datetime.now()
    print(f"分类开始时间: {start_datetime.strftime('%Y-%m-%d %H:%M:%S')}")
    
    if start_idx > 0:
        print(f"从第 {start_idx + 1} 条记录开始继续处理...")
    else:
        print("开始分类处理...")
    
    try:
        for idx in range(start_idx, len(df)):
            row = df.iloc[idx]
            text = str(row['text'])  # 确保是字符串类型
            source_label = row.get('label', '')  # 获取原始标签，如果不存在则为空
            
            try:
                # 进行分类
                prediction = classifier(text)
                
                # 提取预测结果
                predicted_label = prediction[0]['label']
                predicted_score = prediction[0]['score']
                
                # 添加到批量结果列表
                batch_results.append({
                    'text': text,
                    'source_label': source_label,
                    'classified_label': predicted_label,
                    'classified_score': predicted_score
                })
                
                total_processed = idx
                
                # 每处理batch_size条记录或到达最后一条记录时保存
                if len(batch_results) >= batch_size or idx == len(df) - 1:
                    # 保存批量结果
                    mode = 'w' if idx < batch_size and start_idx == 0 else 'a'
                    if save_batch_results(batch_results, output_file, mode):
                        print(f"已处理并保存 {idx + 1}/{len(df)} 条数据")
                        
                        # 保存断点
                        save_checkpoint(checkpoint_file, idx, len(df))
                        
                        # 清空批量结果列表
                        batch_results = []
                    else:
                        print(f"保存第 {idx + 1} 批数据失败，停止处理")
                        break
                
                # 打印进度（每10条显示一次）
                elif (idx + 1) % 10 == 0:
                    print(f"已处理 {idx + 1}/{len(df)} 条数据")
                    
            except Exception as e:
                print(f"处理第 {idx + 1} 行时出错: {e}")
                # 添加错误记录
                batch_results.append({
                    'text': text,
                    'source_label': source_label,
                    'classified_label': 'ERROR',
                    'classified_score': 0.0
                })
                total_processed = idx
                
    except KeyboardInterrupt:
        print(f"\n用户中断处理，已处理到第 {total_processed + 1} 条记录")
        # 保存剩余的批量结果
        if batch_results:
            mode = 'a' if total_processed > 0 else 'w'
            save_batch_results(batch_results, output_file, mode)
        # 保存断点
        save_checkpoint(checkpoint_file, total_processed, len(df))
        return
        
    except Exception as e:
        print(f"\n处理过程中发生错误: {e}")
        # 保存剩余的批量结果
        if batch_results:
            mode = 'a' if total_processed > 0 else 'w'
            save_batch_results(batch_results, output_file, mode)
        # 保存断点
        save_checkpoint(checkpoint_file, total_processed, len(df))
        return
    
    # 记录分类结束时间并计算性能统计
    end_time = time.time()
    end_datetime = datetime.now()
    total_time = end_time - start_time
    processed_count = total_processed - start_idx + 1
    
    print(f"\n分类结束时间: {end_datetime.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"本次处理耗时: {total_time:.2f} 秒")
    print(f"本次处理记录数: {processed_count}")
    if processed_count > 0:
        print(f"平均每条数据耗时: {total_time/processed_count:.3f} 秒")
        print(f"处理速度: {processed_count/total_time:.2f} 条/秒")
    
    # 显示最终统计信息
    try:
        if os.path.exists(output_file):
            results_df = pd.read_csv(output_file)
            print(f"\n总共完成 {len(results_df)} 条记录的分类")
            
            # 打印统计信息
            print("\n分类结果统计:")
            print(results_df['classified_label'].value_counts())
            
            # 显示前几行结果作为示例
            print(f"\n前5行结果预览:")
            print(results_df.head())
            
            # 如果全部完成，删除断点文件
            if len(results_df) == len(df):
                if os.path.exists(checkpoint_file):
                    os.remove(checkpoint_file)
                    print(f"\n所有记录处理完成，已删除断点文件")
        
    except Exception as e:
        print(f"读取最终结果失败: {e}")

if __name__ == "__main__":
    main()