'CohereModel' object has no attribute '_prune_heads'

[mnauf]

System Info

class CoherePreTrainedModel docstring says that "This model inherits from [PreTrainedModel]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.)". However, it doesn't have prune_heads() methods, even though it explicitly mentions that it has in the docstring.

@ArthurZucker

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

from transformers import AutoTokenizer, AutoModel

model_id = "CohereForAI/aya-23-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16)
model.prune_heads({1: [1]})

Expected behavior

I expect the method to prune heads of the cohere model.

[mnauf]

#29622

[mnauf]

@saurabhdash2512 Please have a look

[LysandreJik]

Thanks for your issue @mnauf ! The documentation should be updated, as indeed prune heads is only supported for older models as newer models haven't received this exploration from a research perspective.

[modifiededition]

Hey, I’d like to fix this issue. @LysandreJik , you mentioned that the documentation needs updating—could you clarify which specific documentation so I can work on fixing it?

[jackyjinjing]

Can I support prune_heads feature in this model? @LysandreJik @ArthurZucker

[LysandreJik]

This has been addressed in #33253, so I'm closing the issue - thanks all for offering your help!