Hugging Face's logo

Ryanlal
/
InvarRAG

Model card Files
xet
 Community
InvarRAG
File size: 4,411 Bytes
258adb2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

echo "================================================"
echo "Invar-RAG Query Variants 生成示例"
echo "================================================"

# 配置
DEEPSEEK_API_KEY="${DEEPSEEK_API_KEY:-your_api_key_here}"
BASE_DIR="/Users/dwayneryan/Downloads/RAG-main/opencompass"
DATA_DIR="$BASE_DIR/data/my_datasets"
OUTPUT_DIR="$BASE_DIR/data/variants"
CACHE_DIR="$BASE_DIR/cache/query_variants"

# 检查API key
if [ "$DEEPSEEK_API_KEY" = "your_api_key_here" ]; then
    echo "错误: 请设置DEEPSEEK_API_KEY环境变量"
    echo "export DEEPSEEK_API_KEY='your_actual_key'"
    exit 1
fi

# 创建输出目录
mkdir -p "$OUTPUT_DIR"
mkdir -p "$CACHE_DIR"

echo ""
echo "示例 1: 处理100个样本（测试）"
echo "----------------------------------------"

python "$BASE_DIR/scripts/prepare_invarrag_data.py" convert_nq \
    --input "$DATA_DIR/nq/nq-test.qa.csv" \
    --output "$OUTPUT_DIR/nq-test.json" \
    --max_samples 100

python "$BASE_DIR/scripts/generate_query_variants_deepseek.py" \
    --api_key "$DEEPSEEK_API_KEY" \
    --input "$OUTPUT_DIR/nq-test.json" \
    --output "$OUTPUT_DIR/nq-test-sample-with-variants.json" \
    --num_variants 3 \
    --max_samples 100 \
    --cache_dir "$CACHE_DIR/nq"

echo "✓ 示例 1 完成"

echo ""
echo "示例 2: 处理完整NQ训练集"
echo "----------------------------------------"

python "$BASE_DIR/scripts/prepare_invarrag_data.py" convert_nq \
    --input "$DATA_DIR/nq/nq-train.qa.csv" \
    --output "$OUTPUT_DIR/nq-train.json"

python "$BASE_DIR/scripts/generate_query_variants_deepseek.py" \
    --api_key "$DEEPSEEK_API_KEY" \
    --input "$OUTPUT_DIR/nq-train.json" \
    --output "$OUTPUT_DIR/nq-train-with-variants.json" \
    --num_variants 3 \
    --cache_dir "$CACHE_DIR/nq"

echo "✓ 示例 2 完成"

echo ""
echo "示例 3: 批量处理所有数据集"
echo "----------------------------------------"

DATASETS=(
    "nq:nq-train.qa.csv:nq-train.json:nq-train-with-variants.json:convert_nq"
    "tqa:trivia-train.qa.csv:tqa-train.json:tqa-train-with-variants.json:convert_triviaqa"
)

for dataset in "${DATASETS[@]}"; do
    IFS=':' read -r name input_file json_file output_file convert_cmd <<< "$dataset"
    
    echo ""
    echo "处理: $name"
    
    python "$BASE_DIR/scripts/prepare_invarrag_data.py" $convert_cmd \
        --input "$DATA_DIR/$name/$input_file" \
        --output "$OUTPUT_DIR/$json_file"

    python "$BASE_DIR/scripts/generate_query_variants_deepseek.py" \
        --api_key "$DEEPSEEK_API_KEY" \
        --input "$OUTPUT_DIR/$json_file" \
        --output "$OUTPUT_DIR/$output_file" \
        --num_variants 3 \
        --cache_dir "$CACHE_DIR/$name"
    
    echo "✓ $name 完成"
done

echo ""
echo "验证生成的数据"
echo "----------------------------------------"

python - << EOF
import json
import os

output_dir = "$OUTPUT_DIR"
files = [f for f in os.listdir(output_dir) if f.endswith('.json')]

print(f"生成的文件: {len(files)}")
print("")

for file in files:
    path = os.path.join(output_dir, file)
    with open(path) as f:
        data = json.load(f)
    
    num_samples = len(data)
    num_with_variants = sum(1 for d in data if 'query_variants' in d and d['query_variants'])
    avg_variants = sum(len(d.get('query_variants', [])) for d in data) / num_samples if num_samples > 0 else 0
    
    print(f"文件: {file}")
    print(f"  样本数: {num_samples}")
    print(f"  包含variants: {num_with_variants} ({num_with_variants/num_samples*100:.1f}%)")
    print(f"  平均variants数: {avg_variants:.2f}")
    print("")
EOF

echo ""
echo "缓存统计"
echo "----------------------------------------"

total_cache_size=$(du -sh "$CACHE_DIR" | cut -f1)
cache_files=$(find "$CACHE_DIR" -name "*.json" | wc -l)

echo "缓存目录: $CACHE_DIR"
echo "缓存文件数: $cache_files"
echo "缓存大小: $total_cache_size"

echo ""
echo "================================================"
echo "✓ 所有示例完成"
echo "================================================"
echo ""
echo "生成的文件位于: $OUTPUT_DIR"
echo "缓存文件位于: $CACHE_DIR"
echo ""
echo "下一步："
echo "1. 检查生成的variants质量"
echo "2. 使用variants数据训练模型:"
echo "   python invarrag_train_retriever.py \\"
echo "       --train_data $OUTPUT_DIR/nq-train-with-variants.json \\"
echo "       --corpus data/corpus.jsonl \\"
echo "       ..."
echo ""