python NLP数据集分割大文件

NLP数据文件有时候特别大的文件，需要分割成N个小文件来处理
部分提取：可以提取N份，每份K行
全部分割：分割整个文件，每一份K行

import osdef split_file(filename, outdir,num_lines):""" 将文件按行数进行分割 \nfilename 文件名 \nnum_lines 每份包含的行数 \n"""file_name_without_path_and_extension = os.path.splitext(os.path.basename(filename))[0]with open(filename, 'r') as f:current_chunk = 1current_line = 0current_output = open(f"{outdir}{file_name_without_path_and_extension}{current_chunk}.txt", 'w')for line in f:current_output.write(line)current_line += 1if current_line >= num_lines:current_output.close()current_chunk += 1current_line = 0current_output = open(f"{outdir}{file_name_without_path_and_extension}{current_chunk}.txt", 'w')current_output.close()def split_file_max_chunks(filename,outdir, num_lines, max_chunks):""" 将文件按行数进行分割 \nfilename 文件名 \nnum_lines 每份包含的行数 \nmax_chunks 最大分出多少份 \n"""file_name_without_path_and_extension = os.path.splitext(os.path.basename(filename))[0]with open(filename, 'r') as f:current_chunk = 1current_line = 0current_output = open(f"{outdir}{file_name_without_path_and_extension}{current_chunk}.txt", 'w')for line in f:current_output.write(line)current_line += 1if current_line >= num_lines:current_output.close()if current_chunk >= max_chunks:breakcurrent_chunk += 1current_line = 0current_output = open(f"{outdir}{file_name_without_path_and_extension}{current_chunk}.txt", 'w')  # 这里更新了current_outputcurrent_output.close()def main():large_filename = "./data/large_file_1G.txt"outdir="./docs/"num_lines = 1000  # 每份包含 1000 行split_file(large_filename,outdir, num_lines)# max_chunks = 30   # 最大分出 30 份# split_file_max_chunks(large_filename,outdir, num_lines, max_chunks)if __name__ == "__main__":main()