Error when running layers/test_layer.py

[Yiningww]

At line 44:

hidden_states, outputs = model.generate(**{'input_ids':inputs.input_ids, 'max_new_tokens':64})

There's an error: ValueError: not enough values to unpack (expected 2, got 1). I guess it means that model.generate( ) will only have one output, which is outputs, and can not get hidden_states. I'm wondering how to fix this issue. Thank you very much!

[himanshubeniwal]

+1 on the question above.

[Mongjin]

Hello, I encountered the same issue, but I was able to solve it by revising the Prompting function as shown below.
I believe it is caused by a different version of the transformers library.
In addition, I had to modify the code that decodes the hidden layer’s logits into tokens.

Sorry for the hard-coded solution, but I hope it will be helpful! @Yiningww @himanshubeniwal

If you think that my code may cause an error (i.e., be incorrect), please let me know!

def Prompting(model, prompt, candidate_premature_layers):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(inputs.input_ids, max_new_tokens=64, return_dict_in_generate=True, output_hidden_states=True)
    # hidden_states, outputs = model.generate(**{'input_ids':inputs.input_ids})
    hidden_states = outputs.hidden_states[1:][1:]
    hidden_embed = {}
    hidden_embed_token_level = {}
    lm_head = model.get_output_embeddings()
    for i, early_exit_layer in enumerate(candidate_premature_layers):
        layer_hidden_state = hidden_states[:][early_exit_layer]
        layer_hidden_state = torch.cat(layer_hidden_state, dim=1)
        layer_hidden_logits = lm_head(layer_hidden_state)
        layer_hidden_unembedding = torch.nn.functional.softmax(layer_hidden_logits, dim=-1)
        values, indicies = torch.topk(layer_hidden_unembedding, 100)
        indicies = indicies[0][0][:]
        hidden_embed[early_exit_layer] = tokenizer.decode(indicies)
        hidden_embed_token_level[early_exit_layer] = [tokenizer.decode(tok) for tok in indicies]
    answer = tokenizer.decode(outputs.sequences.tolist()[0]).replace('<pad> ', '')
    answer = answer.replace('</s>', '')
    
    return hidden_embed, hidden_embed_token_level, answer