[Python] Extracting hidden_states in LLaVa's transformer

I need to use LLaVa to obtain the embedding of an image + the query that has been provided to the LLM.

From my understanding, I need the output of the LLaVa model before it goes through the encoding layer.

From this post on StackOverflow, I tried to use pre-hook to save the input of each layer:

def capture_hidden_states(module, module_input):
print(f'module_input: {module_input}\n')

model_path = "liuhaotian/llava-v1.6-vicuna-7b"
model_name = get_model_name_from_path(model_path)

tokenizer, model, image_processor, context_len = load_pretrained_model(
model_path = model_path,
model_base = None,
model_name = model_name,
load_4bit = True, # Save GPU memory
device = "cuda:0" # default
)
model.register_forward_pre_hook(capture_hidden_states)


But my output shows an empty input for every level:

module_input: ()


I see that Hugging Face's version has a output_hidden_states parameter in the forward method, but I am currently unable to make this version of LLaVa run, so I am trying to get this working on the previous code I wrote.

Any help is greatly appreciated.

Continue reading...