| import csv |
| import os |
| import json |
| import argparse |
| import random |
| from process_pretrain_data import get_kmer_sentence |
|
|
|
|
| max_length = 0 |
|
|
| def Process_pair(args): |
| random.seed(42) |
|
|
| root_path = args.file_path.split('/')[-1] |
| train_seq1_file = open(args.file_path+"/"+root_path+"_enhancer.fasta", "r") |
| train_seq2_file = open(args.file_path+"/"+root_path+"_promoter.fasta", "r") |
| train_label_file = open(args.file_path+"/"+root_path+"_label.txt", "r") |
| test_seq1_file = open(args.file_path+"/"+root_path+"_enhancer_test.fasta", "r") |
| test_seq2_file = open(args.file_path+"/"+root_path+"_promoter_test.fasta", "r") |
| test_label_file = open(args.file_path+"/"+root_path+"_label_test.txt", "r") |
|
|
| train_seq1 = train_seq1_file.readlines() |
| train_seq2 = train_seq2_file.readlines() |
| train_label = train_label_file.readlines() |
| test_seq1 = test_seq1_file.readlines() |
| test_seq2 = test_seq2_file.readlines() |
| test_label = test_label_file.readlines() |
|
|
| train_lines = [] |
| test_lines = [] |
| for i in range(len(train_label)): |
| train_lines.append([train_seq1[2*i+1], train_seq2[2*i+1], train_label[i]]) |
| for i in range(len(test_label)): |
| test_lines.append([test_seq1[2*i+1], test_seq2[2*i+1], test_label[i]]) |
|
|
| random.shuffle(train_lines) |
|
|
| if args.dev: |
| num_dev = int(len(train_lines)/10) |
| dev_lines = train_lines[:num_dev] |
| train_lines = train_lines[num_dev:] |
| |
| output_path = make_path(args) |
|
|
| suffix = '.csv' if args.csv else '.tsv' |
| delimiter = ',' if args.csv else '\t' |
|
|
| f_train = open(os.path.join(output_path, "train" + suffix), 'wt') |
| train_w = csv.writer(f_train, delimiter=delimiter) |
| train_w.writerow(["seq1", "seq2", "label"]) |
| if args.dev: |
| f_dev = open(os.path.join(output_path, "dev" + suffix), 'wt') |
| dev_w = csv.writer(f_dev, delimiter=delimiter) |
| dev_w.writerow(["seq1", "seq2", "label"]) |
| os.makedirs(os.path.join(output_path, "test")) |
| f_test = open(os.path.join(output_path, "test", "dev" + suffix), 'wt') |
| test_w = csv.writer(f_test, delimiter=delimiter) |
| test_w.writerow(["seq1", "seq2", "label"]) |
| else: |
| f_test = open(os.path.join(output_path, "dev" + suffix), 'wt') |
| test_w = csv.writer(f_test, delimiter=delimiter) |
| test_w.writerow(["seq1", "seq2", "label"]) |
|
|
| def write_file_pair(lines, writer, seq1_index=0, seq2_index=1, label_index=2): |
| for line in lines: |
| seq1 = get_kmer_sentence(line[seq1_index], kmer=args.kmer, stride=args.stride) |
| seq2 = get_kmer_sentence(line[seq2_index], kmer=args.kmer, stride=args.stride) |
| writer.writerow([seq1, seq2, str(int(line[label_index]))]) |
|
|
| write_file_pair(train_lines, train_w) |
| write_file_pair(test_lines, test_w) |
| |
| if args.dev: |
| write_file_pair(dev_lines, dev_w) |
| |
|
|
| def make_path(args): |
| output_path = args.output_path if args.output_path else os.path.join(args.file_path, str(args.kmer)) |
| if not os.path.exists(output_path): |
| os.makedirs(output_path) |
| return output_path |
|
|
| def write_file(lines, writer, seq_index=2, label_index=3, kmer=6, stride=1): |
| global max_length |
| for line in lines: |
| sentence = get_kmer_sentence(line[seq_index], kmer=kmer, stride=stride) |
| if len(sentence.split()) > max_length: |
| max_length = len(sentence.split()) |
| if label_index == -100: |
| writer.writerow([sentence, str(0)]) |
| else: |
| writer.writerow([sentence, str(line[label_index])]) |
|
|
| def Process(args): |
| random.seed(24) |
|
|
| train = os.path.join(args.file_path, "train.csv") |
| test = os.path.join(args.file_path, "test.csv") |
| train_file = open(train, "r", encoding="utf-8-sig") |
| test_file = open(test, "r", encoding="utf-8-sig") |
|
|
| train_lines = list(csv.reader(train_file, delimiter=",", quotechar=None))[1:] |
| test_lines = list(csv.reader(test_file, delimiter=",", quotechar=None))[1:] |
|
|
| random.shuffle(train_lines) |
| random.shuffle(test_lines) |
|
|
| if args.dev: |
| num_dev = int(len(train_lines)/9) |
| dev_lines = train_lines[:num_dev] |
| train_lines = train_lines[num_dev:] |
|
|
| print(train_lines[0]) |
|
|
| output_path = make_path(args) |
|
|
| suffix = '.csv' if args.csv else '.tsv' |
| delimiter = ',' if args.csv else '\t' |
|
|
|
|
| f_train = open(os.path.join(output_path, "train"+suffix), 'wt') |
| train_w = csv.writer(f_train, delimiter=delimiter) |
| train_w.writerow(["sentence", "label"]) |
| if args.dev: |
| f_dev = open(os.path.join(output_path, "dev"+suffix), 'wt') |
| dev_w = csv.writer(f_dev, delimiter=delimiter) |
| dev_w.writerow(["sentence", "label"]) |
| f_test = open(os.path.join(output_path, "test"+suffix), 'wt') |
| test_w = csv.writer(f_test, delimiter=delimiter) |
| test_w.writerow(["sentence", "label"]) |
| else: |
| f_test = open(os.path.join(output_path, "dev"+suffix), 'wt') |
| test_w = csv.writer(f_test, delimiter=delimiter) |
| test_w.writerow(["sentence", "label"]) |
| |
|
|
| write_file(train_lines, train_w, args.seq_index, args.label_index) |
| write_file(test_lines, test_w, args.seq_index, args.label_index) |
| |
| if args.dev: |
| write_file(dev_lines, dev_w) |
| |
|
|
| print("max length: %d" % (max_length)) |
|
|
|
|
| def Process_UCE(args): |
| len_count = {} |
|
|
| line2index = {} |
|
|
| pred_file = open(args.file_path, "r", encoding="utf-8-sig") |
| pred_lines = list(csv.reader(pred_file, delimiter=",", quotechar=None))[1:] |
|
|
| suffix = '.csv' if args.csv else '.tsv' |
| delimiter = ',' if args.csv else '\t' |
|
|
| f_pred = open(os.path.join(args.output_path, "dev"+suffix), 'wt') |
| pred_w = csv.writer(f_pred, delimiter=delimiter) |
| pred_w.writerow(["sentence", "label"]) |
|
|
| index = 1 |
| line_num = 0 |
| for line in pred_lines: |
| len_count[len(line[8])] = len_count.get(len(line[8]), 0) + 1 |
| len_count[len(line[-2])] = len_count.get(len(line[-2]), 0) + 1 |
|
|
| cur_index = [index, index+1] |
| ref = get_kmer_sentence(line[8], args.kmer, args.stride) |
| pred_w.writerow([ref, 0]) |
|
|
| mut1 = get_kmer_sentence(line[-2], args.kmer, args.stride) |
| pred_w.writerow([mut1, 0]) |
|
|
| index += 2 |
|
|
| if line[-2] != line[-1]: |
| len_count[len(line[-1])] = len_count.get(len(line[-1]), 0) + 1 |
| mut2 = get_kmer_sentence(line[-1], args.kmer, args.stride) |
| pred_w.writerow([mut2, 0]) |
| cur_index.append(index) |
| index += 1 |
| |
| line2index[line_num] = cur_index |
| line_num += 1 |
| |
| with open(os.path.join(args.output_path, "line2index.json"), "w") as f: |
| json.dump(line2index, f) |
| with open(os.path.join(args.output_path, "lencount.json"), "w") as f: |
| json.dump(len_count, f) |
|
|
|
|
| def Process_Virus(args): |
| file_path = args.file_path |
|
|
| all_files = os.listdir(file_path) |
| all_files = [f for f in all_files if not f.startswith("unclass")] |
| all_lines = [] |
| for i, f in enumerate(all_files): |
| f_dir = os.path.join(file_path, f) |
| cur_file = open(f_dir, "r", encoding="utf-8-sig") |
| cur_lines = list(csv.reader(cur_file, delimiter=",", quotechar=None))[1:] |
| all_lines.extend(cur_lines) |
| |
|
|
| suffix = '.csv' if args.csv else '.tsv' |
| delimiter = ',' if args.csv else '\t' |
|
|
| f_pred = open(os.path.join(args.output_path, "dev"+suffix), 'wt') |
| pred_w = csv.writer(f_pred, delimiter=delimiter) |
| pred_w.writerow(["sentence", "label"]) |
|
|
| index = 1 |
| line_num = 0 |
| for line in pred_lines: |
| cur_index = [index, index+1] |
| ref = get_kmer_sentence(line[8], args.kmer, args.stride) |
| pred_w.writerow([ref, 0]) |
|
|
| mut1 = get_kmer_sentence(line[-2], args.kmer, args.stride) |
| pred_w.writerow([mut1, 0]) |
|
|
| index += 2 |
|
|
| if line[-2] != line[-1]: |
| len_count[len(line[-1])] = len_count.get(len(line[-1]), 0) + 1 |
| mut2 = get_kmer_sentence(line[-1], args.kmer, args.stride) |
| pred_w.writerow([mut2, 0]) |
| cur_index.append(index) |
| index += 1 |
| |
| line2index[line_num] = cur_index |
| line_num += 1 |
| |
| with open(os.path.join(args.output_path, "line2index.json"), "w") as f: |
| json.dump(line2index, f) |
| with open(os.path.join(args.output_path, "lencount.json"), "w") as f: |
|
|
|
|
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| "--kmer", |
| default=1, |
| type=int, |
| help="K-mer", |
| ) |
| parser.add_argument( |
| "--stride", |
| default=1, |
| type=int, |
| help="stride in getting kmer sequence", |
| ) |
| parser.add_argument( |
| "--file_path", |
| default=None, |
| type=str, |
| help="The path of the file to be processed", |
| ) |
| parser.add_argument( |
| "--output_path", |
| default=None, |
| type=str, |
| help="The path of the processed data", |
| ) |
| parser.add_argument( |
| "--dev", |
| action="store_true", |
| help="Use this flag to split data as (8:1:1), else (9:1)", |
| ) |
| parser.add_argument( |
| "--csv", |
| action="store_true", |
| help="if output csv file or not, if not, output tsv", |
| ) |
| parser.add_argument( |
| "--pair", |
| action="store_true", |
| help="Use this flag to split data as (8:1:1), else (9:1)", |
| ) |
| parser.add_argument( |
| "--uce", |
| action="store_true", |
| help="Use this flag to split data as (8:1:1), else (9:1)", |
| ) |
| parser.add_argument( |
| "--seq_index", |
| default=2, |
| type=int, |
| help="index of seq in the original csv file", |
| ) |
| parser.add_argument( |
| "--label_index", |
| default=3, |
| type=int, |
| help="index of label in the original csv file", |
| ) |
| args = parser.parse_args() |
| |
| if args.pair: |
| Process_pair(args) |
| elif args.uce: |
| Process_UCE(args) |
| else: |
| Process(args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
