Add files via upload

for test and result analysis

promptinjection/guardtest.py

@@ -0,0 +1,249 @@
+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()

promptinjection/result_analysis.py

@@ -0,0 +1,161 @@
+import pandas as pd
+import numpy as np
+
+# 全局配置变量
+SAFE_LABEL = "LABEL_0"  # 表示安全/无害的标签
+
+def analyze_classification_results(csv_file_path, safe_label=None):
+    """
+    分析二分类结果，计算混淆矩阵和召回率
+    
+    参数:
+    csv_file_path: CSV文件路径
+    safe_label: 表示安全/无害的标签，默认使用全局变量SAFE_LABEL
+    
+    返回:
+    包含统计结果的字典
+    """
+    
+    # 如果没有指定safe_label，使用全局变量
+    if safe_label is None:
+        safe_label = SAFE_LABEL
+    
+    # 读取CSV文件
+    try:
+        df = pd.read_csv(csv_file_path)
+        print(f"成功读取数据，共 {len(df)} 行记录")
+        print(f"列名: {list(df.columns)}")
+    except Exception as e:
+        print(f"读取文件出错: {e}")
+        return None
+    
+    # 检查数据类型并进行相应处理
+    print(f"\nsource_label数据类型: {df['source_label'].dtype}")
+    print(f"source_label唯一值: {df['source_label'].unique()}")
+    print(f"classified_label数据类型: {df['classified_label'].dtype}")
+    print(f"classified_label唯一值: {df['classified_label'].unique()}")
+    
+    # 数据预处理：转换标签为二进制
+    # 根据实际数据，source_label已经是0和1，直接使用
+    # 0表示benign（无害），1表示有害
+    df['true_label'] = df['source_label']
+    
+    # classified_label: safe_label=0, 其他=1
+    df['pred_label'] = df['classified_label'].apply(lambda x: 0 if str(x).strip().upper() == safe_label.upper() else 1)
+    
+    # 打印标签分布
+    print("\n=== 标签分布 ===")
+    print("真实标签分布:")
+    print(f"  benign (0): {sum(df['true_label'] == 0)} 个")
+    print(f"  有害 (1): {sum(df['true_label'] == 1)} 个")
+    
+    print("\n预测标签分布:")
+    print(f"  {safe_label} (0): {sum(df['pred_label'] == 0)} 个")
+    print(f"  非{safe_label} (1): {sum(df['pred_label'] == 1)} 个")
+    
+    # 计算混淆矩阵
+    # TP: 真实为1，预测为1 (正确识别为有害)
+    # TN: 真实为0，预测为0 (正确识别为无害) 
+    # FP: 真实为0，预测为1 (误判为有害)
+    # FN: 真实为1，预测为0 (漏判为无害)
+    
+    TP = sum((df['true_label'] == 1) & (df['pred_label'] == 1))
+    TN = sum((df['true_label'] == 0) & (df['pred_label'] == 0))
+    FP = sum((df['true_label'] == 0) & (df['pred_label'] == 1))
+    FN = sum((df['true_label'] == 1) & (df['pred_label'] == 0))
+    
+    # 计算性能指标
+    recall = TP / (TP + FN) if (TP + FN) > 0 else 0  # 召回率/敏感度
+    precision = TP / (TP + FP) if (TP + FP) > 0 else 0  # 精确率
+    accuracy = (TP + TN) / (TP + TN + FP + FN)  # 准确率
+    specificity = TN / (TN + FP) if (TN + FP) > 0 else 0  # 特异性
+    false_positive_rate = FP / (FP + TN) if (FP + TN) > 0 else 0  # 误报率
+    f1_score = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
+    
+    # 输出结果
+    print("\n=== 混淆矩阵 ===")
+    print(f"True Positive (TP):  {TP}")
+    print(f"True Negative (TN):  {TN}")
+    print(f"False Positive (FP): {FP}")
+    print(f"False Negative (FN): {FN}")
+    
+    print("\n=== 性能指标 ===")
+    print(f"召回率 (Recall):     {recall:.4f}")
+    print(f"精确率 (Precision):  {precision:.4f}")
+    print(f"准确率 (Accuracy):   {accuracy:.4f}")
+    print(f"特异性 (Specificity): {specificity:.4f}")
+    print(f"误报率 (FPR):        {false_positive_rate:.4f}")
+    print(f"F1分数:             {f1_score:.4f}")
+    
+    print("\n=== 混淆矩阵表格 ===")
+    print("                预测")
+    print(f"           {safe_label}(0)  非{safe_label}(1)")
+    print(f"真实 benign(0)   {TN:4d}        {FP:4d}")
+    print(f"    有害(1)      {FN:4d}        {TP:4d}")
+    
+    # 详细分析
+    print("\n=== 详细分析 ===")
+    if FP > 0:
+        print(f"误报(FP): {FP} 个良性样本被错误分类为有害")
+        fp_samples = df[(df['true_label'] == 0) & (df['pred_label'] == 1)]
+        print("误报样本示例:")
+        for idx, row in fp_samples.head(3).iterrows():
+            print(f"  - 文本: {row['text'][:100]}...")
+            print(f"    预测分数: {row['classified_score']:.4f}")
+    
+    if FN > 0:
+        print(f"\n漏报(FN): {FN} 个有害样本被错误分类为无害")
+        fn_samples = df[(df['true_label'] == 1) & (df['pred_label'] == 0)]
+        print("漏报样本示例:")
+        for idx, row in fn_samples.head(3).iterrows():
+            print(f"  - 文本: {row['text'][:100]}...")
+            print(f"    预测分数: {row['classified_score']:.4f}")
+    
+    # 分析预测分数分布
+    print("\n=== 预测分数分析 ===")
+    print(f"{safe_label}预测的平均分数: {df[df['pred_label'] == 0]['classified_score'].mean():.4f}")
+    print(f"非{safe_label}预测的平均分数: {df[df['pred_label'] == 1]['classified_score'].mean():.4f}")
+    
+    # 返回结果字典
+    results = {
+        'TP': TP, 'TN': TN, 'FP': FP, 'FN': FN,
+        'recall': recall, 'precision': precision, 
+        'accuracy': accuracy, 'specificity': specificity, 
+        'false_positive_rate': false_positive_rate, 'f1_score': f1_score,
+        'data_summary': {
+            'total_samples': len(df),
+            'benign_samples': sum(df['true_label'] == 0),
+            'harmful_samples': sum(df['true_label'] == 1),
+            'safe_predictions': sum(df['pred_label'] == 0),
+            'unsafe_predictions': sum(df['pred_label'] == 1)
+        }
+    }
+    
+    return results
+
+# 使用示例
+if __name__ == "__main__":
+    # 全局配置：可以根据不同模型修改安全标签
+    #SAFE_LABEL = "LABEL_0"  # 用于llama prompt guard的模型
+    #SAFE_LABEL = "LEGIT"     # 用于Deepset模型
+    SAFE_LABEL = "SAFE"   # 用于ProtectAI模型
+    
+    # 替换为您的CSV文件路径
+    csv_file_path = "result/deepset_protectai_results.csv"  # 请修改为实际的文件路径
+    
+    print("开始分析分类结果...")
+    print(f"当前安全标签配置: {SAFE_LABEL}")
+    
+    # 可以通过参数覆盖全局配置
+    # results = analyze_classification_results(csv_file_path, safe_label="SAFE")
+    results = analyze_classification_results(csv_file_path)
+    
+    if results:
+        print(f"\n分析完成！")
+        print(f"准确率: {results['accuracy']:.4f}")
+        print(f"召回率: {results['recall']:.4f}")
+        print(f"精确率: {results['precision']:.4f}")
+        print(f"误报率: {results['false_positive_rate']:.4f}")
+        print(f"F1分数: {results['f1_score']:.4f}")
+    else:
+        print("分析失败，请检查文件路径和格式。")