Upload data3/load_dataset.py with huggingface_hub

data3/load_dataset.py

@@ -0,0 +1,75 @@
+import csv
+
+csv.field_size_limit(10 * 1024 * 1024 * 1024)  # 10MB
+
+# length_max = 300
+# length_max = 28906320
+# length_max = 8290
+
+# with open('res.csv', 'r', encoding='utf-8') as f:
+#     reader = csv.reader(f)
+#     id_set = set()
+#     for i, row in enumerate(reader):
+#         id_set.add(row[0])
+#     print(len(id_set))
+
+with open('res1.csv', 'r', encoding='utf-8') as f:
+    reader = csv.reader(f)
+    id_set = set()
+    for i, row in enumerate(reader):
+        id_set.add(row[0])
+    print(len(id_set))
+
+with open('res2.csv', 'r', encoding='utf-8') as f:
+    reader = csv.reader(f)
+    for i, row in enumerate(reader):
+        if row[0] in id_set:
+            id_set.remove(row[0])
+    print(len(id_set))
+
+length_max = len(id_set)
+
+def load_dataset():
+    prompt = open('score_prompt.txt', 'r', encoding='utf-8').read()
+    # prompt = open('is_sci_prompt.txt', 'r', encoding='utf-8').read()
+    with open('/home/weifengsun/tangou1/domain_code/src/datasets/data_merged/dataset_all.csv', 'r', encoding='utf-8') as f:
+        reader = csv.reader(f)
+        amount = 0
+        # ['', 'text', 'repo_name', 'path', 'language', 'license', 'size', 'keyword', 'text_hash', 'config', 'split', 'repo_path', 'ds_source']
+        for i, row in enumerate(reader):
+            if i == 0:
+                continue
+            if amount > length_max:
+                return
+            if len(row[1]) > 100000 or len(row[1]) < 2000:
+                continue
+            if row[0] in id_set:
+                amount += 1
+                yield prompt.format(CODE_FILE=add_line_numbers(row[1])), row[0]
+                # yield prompt.format(CODE_FILE=row[1][:20000]), row[0]
+
+
+def add_line_numbers(text: str) -> str:
+    out_lines = []
+    line_num = 1
+    start = 0
+    for line in text.splitlines(keepends=True):
+        out_lines.append(f"{line_num} " + line)
+        line_num += 1
+    return "".join(out_lines)
+
+
+if __name__ == "__main__":
+    # with open('/home/weifengsun/tangou1/domain_code/src/datasets/data_merged/dataset_all.csv', 'r', encoding='utf-8') as f:
+    #     reader = csv.reader(f)
+    #     # ['', 'text', 'repo_name', 'path', 'language', 'license', 'size', 'keyword', 'text_hash', 'config', 'split', 'repo_path', 'ds_source']
+    #     for i, row in enumerate(reader):
+    #         if i > 100:
+    #             break   
+    #         print(row[0])
+    # amount = 0
+    # for i in load_dataset():
+    #     amount += 1
+    # print(amount)
+    # print(len(id_set))
+    pass