def get_ranged_completions()

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

• 10 lines of code
• 1 McCabe index (conditional complexity)

def get_ranged_completions(dataset_name, entity_name, indices_path):
    with open(f"./data/{dataset_name}/{entity_name}/completions.json", "r") as f:
        entity_data = json.load(f)
    dataset_embeddings, entity_embeddings = load_embeddings(dataset_name, entity_name, indices_path)
    input_dim, output_dim = dataset_embeddings.shape[-1], entity_embeddings.shape[-1]
    model = load_model(input_dim, hidden_dim, output_dim, entity_name)
    ffn_embeddings = get_answer_embedding(model, dataset_embeddings, output_dim)
    ranked_completions = get_closest_embedding(entity_data, entity_embeddings, ffn_embeddings)
    with open(f"./data/{dataset_name}/{entity_name}/ranked_completions.json", 'w') as f:
        json.dump(ranked_completions, f, indent=2)