def make_jsonl_samples_iter()

in summarize_from_feedback/task_data.py [0:0]

• 17 lines of code
• 5 McCabe index (conditional complexity)

def make_jsonl_samples_iter(input_path, layout: Optional[ModelLayout] = None):
    """
    Makes an iterator reading examples out of all the samples.[0-9]*.jsonl files in the given path,
    distributed across replicas according to the layout.
    """
    if blobs.is_blob_url(input_path):
        local_input_dir = blobs.download_directory_cached(input_path)
    else:
        local_input_dir = input_path
    input_file_names = glob(os.path.join(local_input_dir, "samples.[0-9]*.jsonl"))

    def all_examples():
        for file_name in input_file_names:
            with bf.BlobFile(file_name, "r") as f:
                for line in f:
                    encoded_example = json.loads(line)
                    example = jsonl_encoding.decode_example(encoded_example)
                    yield example

    d = all_examples()
    if layout:
        d = even_more_itertools.distribute(d, layout)
    return d