<commit_before>import pandas as pd
import sys
import argparse
import os

parser = argparse.ArgumentParser(description='Remove columns in segmentation file that fail QC.')
parser.add_argument('-s','--segmentation-file', metavar='SEGMENTATION_FILE', required=True, help='A file containing the segmentation of all samples')
parser.add_argument('-g','--keep-good', action='store_true', help='Keep only the listed samples (default: discard listed samples)')
parser.add_argument('-n','--sample-names', metavar='SAMPLE_NAME', required=True, nargs='*', help='Names of the samples to remove')

args = parser.parse_args()

data = pd.read_csv(args.segmentation_file, delim_whitespace=True, index_col=[0,1,2])

index = { os.path.basename(c).split('.')[0] : c for c in data.columns }

bad_samples = [ index[b] for b in args.sample_names ]

if args.keep_good:
    subset = data[bad_samples]
else:
    subset = data.drop(bad_samples, 1)


subset.to_csv(sys.stdout, index=True, sep='\t')

<commit_msg>Added -g flag to select_samples.py<commit_after>import pandas as pd
import sys
import argparse
import os

parser = argparse.ArgumentParser(description='Remove columns in segmentation file that fail QC.')
parser.add_argument('-s','--segmentation-file', metavar='SEGMENTATION_FILE', required=True, help='A file containing the segmentation of all samples')
parser.add_argument('-g','--keep-good', action='store_true', help='Keep only the listed samples (default: discard listed samples)')
parser.add_argument('-n','--sample-names', metavar='SAMPLE_NAME', required=True, nargs='*', help='Names of the samples to remove')

args = parser.parse_args()

if __name__ == '__main__':

    data = pd.read_csv(args.segmentation_file, delim_whitespace=True, index_col=[0,1,2])

    index = { }

    for c in data.columns:
        index[os.path.basename(c).split('.')[0]] = c  

    bad_samples = [ index[b] for b in args.sample_names ]

    if args.keep_good:
        subset = data[bad_samples]
    else:
        subset = data.drop(bad_samples, 1)


    subset.to_csv(sys.stdout, index=True, sep='\t')


{'fraction_non_alphanumeric': 0.1446480231436837, 'fraction_numerical': 0.0048216007714561235, 'mean_word_length': 5.085043988269795, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4356298', 'n_tokens_mistral': 687, 'n_tokens_neox': 663, 'n_words': 169}