diff --git a/code/data_gen.cpp b/code/data_gen.cpp
index 0d00122..d9f821d 100644
--- a/code/data_gen.cpp
+++ b/code/data_gen.cpp
@@ -5,16 +5,17 @@
 #include <map>
 #include <vector>
 #include <string.h>
+#include <cstdlib>
 //#include "json.hpp"
 using namespace std;
 //using json = nlohmann::json;
-const int n = 5;
-const int nn = n*(n-1)/2;
+int n = 5;
+int nn;
 const int max_n = 8;
 //const int total_graph_number = ;
 //const int total_order_number = ;
 
-int power_3[10];
+int power_3[max_n*(max_n-1)/2];
 int flag[1000010];
 bool f[1000010];
 int mp[100][100];
@@ -36,8 +37,8 @@ int unique_DAG_num = 0;
 struct Shuffle_DAG{
     int *order_list;
     int graph_idx;//是下表
-}shuffleDag[1000010];
-int flag_generated_graph[1000010] = {}, number_shuffleDag = 0;// If the graph is already generated, then skip it.
+}shuffleDag[15000010];
+int flag_generated_graph[15000010] = {}, number_shuffleDag = 0;// If the graph is already generated, then skip it.
 //using namespace std;
 namespace d_seperate{
     inline ll read(){
@@ -386,8 +387,8 @@ namespace node_relations{
                 if (i != j) {
                     int new_x = order_list[i], new_y = order_list[j];
                     for (int k = 0; k < 7; k++) {
-                        if ((graph_2.pair_relations[new_x][new_y]>>k)&1)
-                            graph_1.pair_relation_count[i][j][k]++;
+                        if ((graph_2.pair_relations[i][j]>>k)&1)
+                            graph_1.pair_relation_count[new_x][new_y][k]++;
                     }
                 }
     }
@@ -512,15 +513,19 @@ namespace json_output{
         putchar('}');puts("");
     }
     void jump_final_answer(){
-        freopen("/Users/lvzhiheng/causal_relation_new_n=6.jsonl", "w", stdout);
+        char filename[256];
+        sprintf(filename, "../data/causal_relation_new_n=%d.jsonl", n);
+        freopen(filename, "w", stdout);
         for (int i = 1; i <= unique_DAG_num; i++){
             jump_single_graph(uniqueDag[i]);
         }
     }
 }
-int main() {
+int main(int argc, char* argv[]) {
+    if (argc > 1) n = atoi(argv[1]);
+    nn = n*(n-1)/2;
     power_3[0]=1;
-    for (int i=1;i<=9;i++)power_3[i]=power_3[i-1]*3;
+    for (int i=1;i<max_n*(max_n-1)/2;i++)power_3[i]=power_3[i-1]*3;
     node_relations::generate_node_relations();
     json_output::jump_final_answer();
     return 0;
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..08501dd
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,16 @@
+numpy
+pandas
+torch
+transformers
+scikit-learn
+tqdm
+torchmetrics
+matplotlib
+seaborn
+nltk
+openai
+huggingface_hub
+efficiency
+shap
+scipy
+pytorch_transformers