Unable to retrieve hidden_states

[vsoesanto]

I converted a locally saved T5 checkpoint to ONNX using FastT5:

>>> from fastT5 import export_and_get_onnx_model
>>> from transformers import AutoTokenizer

>>> model_checkpoint = "path/to/checkpoint"
>>> model = export_and_get_onnx_model(model_name)

I tested it for inference:

>>> tokenizer = AutoTokenizer.from_pretrained(model_name)

>>> token = tokenizer(input_terms, max_length=512 * 2, padding=True, truncation=True, return_tensors='pt')

>>> out = model.generate(input_ids=token['input_ids'].to('cpu'),
                            attention_mask=token['attention_mask'].to('cpu'),
                            return_dict_in_generate=True,
                            max_length=512 * 2,
                            num_beams=1,
                            output_scores=True,
                            output_hidden_states=True)

>>> out.encoder_hidden_states
>>> out.decoder_hidden_states
(None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
 None,
...

>>> out
GreedySearchEncoderDecoderOutput(sequences=tensor([[  0, 119, 114, 102, 108, 111, 108, 125, 120, 112, 100, 101,  35,  53, ...
...
), , encoder_attentions=None, encoder_hidden_states=None, decoder_attentions=None, cross_attentions=None, decoder_hidden_states=(None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None))

The hidden states are all None.

Is there any way that I can retrieve the hidden states for both encoder and decoder?

[vsoesanto]

@Ki6an I also tried by adding "output_hidden_states" = True in the onnx model's config. Also made sure this argument is added in model.generate() call, but still no luck. Any idea how I can retrieve the encoder/decoder hidden states?

[Ki6an]

sorry for the late reply,

you can get the hidden states of the encoder easily just by sending in the input_ids and attention mask to the encoder as shown below

...
model = export_and_get_onnx_model(model_name)
encoder = model.encoder

hidden_state = encoder(input_ids, attention_mask)

but for the decoder, you need to make lots of changes. you can start by making changes here

fastT5/fastT5/onnx_models_structure.py

Lines 52 to 62 in 20441b3

	def forward(self, input_ids, attention_mask, encoder_hidden_states):
	decoder_output = self.decoder(
	input_ids=input_ids,
	encoder_attention_mask=attention_mask,
	encoder_hidden_states=encoder_hidden_states,
	)
	return (
	self.lm_head(decoder_output[0] * (self.config.d_model ** -0.5)),
	decoder_output[1],
	)

decoder_output[0] is the last_hidden_state

make it return that value as well

     return ( 
         decoder_output[0],
         self.lm_head(decoder_output[0] * (self.config.d_model ** -0.5)), 
         decoder_output[1], 
     ) 

also, do the same changes for the decoder.

then, retrieve those values from ort session here and here

finally, pass those values here

fastT5/fastT5/onnx_models.py

Line 199 in 20441b3

return Seq2SeqLMOutput(logits=logits, past_key_values=past_key_values)

as decoder_hidden_states=