def get_ffn_ranged_entities()

in src/scripts/extract_data.py [0:0]

• 13 lines of code
• 6 McCabe index (conditional complexity)

def get_ffn_ranged_entities(dataset_name, data, indices_path):
    ranged_entities = {}
    keys = create_indices(dataset_name, indices_path)
    for entity_name in Entity:
        dump_ffn_input_completions(entity_name, dataset_name, data, keys)
        get_embeddings(dataset_name, ModelType.ADA, entity_name, indices_path)
        get_ranged_completions(dataset_name, entity_name, indices_path)
        if entity_name == Entity.EVENT_TYPE or entity_name == Entity.OBJECT_TYPE:
            ranged_entities[entity_name] = get_grouped_entities(entity_name, dataset_name, indices_path)
        else:
            with open(f"./data/{dataset_name}/{entity_name}/ranked_completions.json", 'r') as f:
                ranged_entities[entity_name] = json.load(f)
    return {key: {entity_name.value: ranged_entities[entity_name][key] for entity_name in Entity} for key in keys}