improve serialization format of extract_representation

Created by: guillaume-wisniewski

we are serializing everything to columns "embedding_layer_N" and "attention_layer_N_head_M". It would be better to have hierarchical columns and make the number explicit (e.g. to be able to easily select all head in a given layer) rather than having to parse the string and extract the information.