Add evolla rebase main
Author: zhoubay
Pull request: yes
Labels: name: New model

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->
Fix issue #36231 

Evolla is an advanced 80-billion-parameter protein-language generative model designed to decode the molecular language of proteins. It integrates information from protein sequences, structures, and user queries to generate precise and contextually nuanced insights into protein function. Trained on an unprecedented AI-generated dataset of 546 million protein question-answer pairs and 150 billion word tokens, Evolla significantly advances research in proteomics and functional genomics, providing expert-level insights and shedding light on the molecular logic encoded in proteins.

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@amyeroberts @qubvel @ArthurZucker 

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
> Super nice to see a new model for proteins! 🚀 ! @Cyrilvallez can you do a first pass?

Thanks! I'm currently working on the tests, it seems great in my device, but lots of tests are failed :(

Here's the error in `tests_processors`, which is the only test failed:
```
OSError: Can't load tokenizer for 'meta-llama/Meta-Llama-3-8B-Instruct'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the same name. Otherwise, make sure 'meta-llama/Meta-Llama-3-8B-Instruct' is the correct path to a directory containing all relevant files for a LlamaTokenizerFast tokenizer.
```

I believe the reason is the access problem. Do you think there are some way to get off that?

Here I met a problem in `tests_processors`. I'm using two tokenizers for my `EvollaProcessor` class, one is `EsmTokenizer` and another is `LlamaTokenizerFast`. When I tried to save the `EvollaProcessor` class, the two tokenizer collide with each other, since they are both `Tokenizer` and the saved files are named exactly the same way and therefore one tokenizer will replace another one's file. 

What I did was changing a little bit of `save_pretrained` in `EvollaProcessor` class. Here's the snippet in [here](https://github.com/zhoubay/transformers/blob/c8866b79c24b6763002e04b6b5d94fddb27ade1d/src/transformers/models/evolla/processing_evolla.py#L375-L380):

```python
    def save_pretrained(self, save_directory, push_to_hub=False, **kwargs):
        # only save the protein tokenizer in sub_dir
        self.protein_tokenizer.save_pretrained(
            os.path.join(save_directory, self.protein_tokenizer_dir_name), push_to_hub=push_to_hub, **kwargs
        )
        return super().save_pretrained(save_directory, push_to_hub, **kwargs)
```

In short words, I tried to save `protein_tokenizer` in a `sub_dir` and then load it from the `sub_dir`.

To load it, I rewrite the `_get_arguments_from_pretrained`([here](https://github.com/zhoubay/transformers/blob/c8866b79c24b6763002e04b6b5d94fddb27ade1d/src/transformers/models/evolla/processing_evolla.py#L382-L408)) in `load_pretrained` function ([here](https://github.com/zhoubay/transformers/blob/dd16acb8a3e93b643aa374c9fb80749f5235c1a6/src/transformers/processing_utils.py#L1070)).

My `_get_arguments_from_pretrained`:
```python
    @classmethod
    # reimplement this to add sub_dir_name support
    def _get_arguments_from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        args = []
        for attribute_name in cls.attributes:
            class_name = getattr(cls, f"{attribute_name}_class")
            sub_dir_name = getattr(cls, f"{attribute_name}_dir_name", None)
            if isinstance(class_name, tuple):
                classes = tuple(getattr(transformers_module, n) if n is not None else None for n in class_name)
                use_fast = kwargs.get("use_fast", True)
                if use_fast and classes[1] is not None:
                    attribute_class = classes[1]
                else:
                    attribute_class = classes[0]
            else:
                attribute_class = getattr(transformers_module, class_name)

            if sub_dir_name is not None:
                args.append(
                    attribute_class.from_pretrained(
                        os.path.join(pretrained_model_name_or_path, sub_dir_name), **kwargs
                    )
                )
            else:
                args.append(attribute_class.from_pretrained(pretrained_model_name_or_path, **kwargs))

        return args
```

The added lines are the following.

```python
            if sub_dir_name is not None:
                args.append(
                    attribute_class.from_pretrained(
                        os.path.join(pretrained_model_name_or_path, sub_dir_name), **kwargs
                    )
                )
            else:
                args.append(attribute_class.from_pretrained(pretrained_model_name_or_path, **kwargs))
----------------------------------------------------------------------------------------------------------------------------------
            args.append(attribute_class.from_pretrained(pretrained_model_name_or_path, **kwargs))
```

And the `_get_arguments_from_pretrained` functions in `from_pretrained` as [here](https://github.com/zhoubay/transformers/blob/dd16acb8a3e93b643aa374c9fb80749f5235c1a6/src/transformers/processing_utils.py#L1070) (I removed the docstring to make it short):

```python
    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, os.PathLike],
        cache_dir: Optional[Union[str, os.PathLike]] = None,
        force_download: bool = False,
        local_files_only: bool = False,
        token: Optional[Union[str, bool]] = None,
        revision: str = "main",
        **kwargs,
    ):
        kwargs["cache_dir"] = cache_dir
        kwargs["force_download"] = force_download
        kwargs["local_files_only"] = local_files_only
        kwargs["revision"] = revision

        use_auth_token = kwargs.pop("use_auth_token", None)
        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            token = use_auth_token

        if token is not None:
            kwargs["token"] = token

        args = cls._get_arguments_from_pretrained(pretrained_model_name_or_path, **kwargs)
        processor_dict, kwargs = cls.get_processor_dict(pretrained_model_name_or_path, **kwargs)

        return cls.from_args_and_dict(args, processor_dict, **kwargs)
```

However, the problem is, it works very well in local file system, but when it comes to loading models from the hub, the error comes as:

```
OSError: Incorrect path_or_model_id: 'westlake-repl/Evolla-10B-hf/protein_tokenizer'. Please provide either the path to a local folder or the repo_id of a model on the Hub.
```

I'm totally aware of the problem, but how to solve this kind of issue? Do you have any other solutions to my case?

Thanks in advance!
@ArthurZucker @Cyrilvallez  



Hey! Actually, cc @zucchini-nlp, do you know if we have precedent for processors using 2 tokenizers? 🤔

In the meantime, @zhoubay I see that quite a lot of code in `modeling_evolla.py` is adapted from Llama, and in `protein.py` from CLIP. We now use [modular](https://huggingface.co/docs/transformers/main/en/modular_transformers) to handle such cases! Basically, you only write your model definition in the modular file using inheritance from other models, and then our linter will create the modeling files 🤗 You can see e.g. [qwen2](https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2/modular_qwen2.py) for a first simple example, and [qwen2.5-vl](https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py) for a more advanced example closer to yours 🤗 Note that if needed, you should be able to create 2 modulars, for the protein model and for the evolla one, if it makes sense to separate them 👌

Let me know if you need any further help concerning modular! 🤗

@Cyrilvallez hey! yeah, we have InstructBlip with 2 tokenizers, one for inputs and another for Qformer module. We added a [custom from/save pretrained](https://github.com/huggingface/transformers/blob/9b479a245b793cac2a8b2e87c6d8e81bb24e20c4/src/transformers/models/instructblip/processing_instructblip.py#L191-L221) to get the tokenizers. On the hub, the other tokenizer is saved in a separate file called `qformer_tokenizer`

Thank you @zucchini-nlp and @Cyrilvallez ! 

I've solved the problem by refering the code from `instructblip`! However, even I pass the `test_processor`, I still failed in `tests_generate` because of `tests/models/video_llava/test_modeling_video_llava.py::VideoLlavaForConditionalGenerationModelTest::test_assisted_decoding_with_logits_to_keep` which is passed in my environment.

What should I do to progress the process?

As for the modular issue @Cyrilvallez mentioned, is it compulsory to change it to that version? I'm reading the docs, and very favor to apply it to reduce my reused codes without hurting the performance! 

Feel free to ignore tests from other models. The test failing for you was fixed yesterday, so rebasing might help. If there are other flaky tests with vision LLMs, let me know, I am working on decreasing flakiness to minimum :)

@zucchini-nlp Thanks! As for the `instructblip`, I've noticed a little flakiness with their implementing.

For the [referenced lines](https://github.com/huggingface/transformers/blob/9b479a245b793cac2a8b2e87c6d8e81bb24e20c4/src/transformers/models/instructblip/processing_instructblip.py#L200-L207), they removed `qformer_tokenizer` and then append it as the last item of `self.attributes`. However, this would work only because the `qformer_tokenizer` is originally the last of the `self.attributes`, see [here](https://github.com/huggingface/transformers/blob/9b479a245b793cac2a8b2e87c6d8e81bb24e20c4/src/transformers/models/instructblip/processing_instructblip.py#L74).

```python
# We modify the attributes so that only the tokenizer and image processor are saved in the main folder
qformer_present = "qformer_tokenizer" in self.attributes
if qformer_present:
    self.attributes.remove("qformer_tokenizer")

outputs = super().save_pretrained(save_directory, **kwargs)

if qformer_present:
    self.attributes += ["qformer_tokenizer"]
```

This might cause unexpected error when you change the order of the `self.attributes` to, say, `attributes = ["image_processor", "qformer_tokenizer", "tokenizer"]`, and at the same time change the signiture of `__init__` into `def __init__(self, image_processor, qformer_tokenizer, tokenizer, num_query_tokens=None, **kwargs):` (since the `attributes` should have same order of the `__init__` function). If you change that way, there'll be an unexpected error. 

My proposal is as my recent commit, add a variable to track where the popped out attribute is, as shown below:

```python
# we modify the attributes so that only the text tokenizer are saved in the main folder
protein_tokenizer_present = "protein_tokenizer" in self.attributes
# find the correct position of it in the attributes list
protein_tokenizer_index = self.attributes.index("protein_tokenizer") if protein_tokenizer_present else None
if protein_tokenizer_present and protein_tokenizer_index is not None:
    self.attributes.remove("protein_tokenizer")

outputs = super().save_pretrained(save_directory, **kwargs)

if protein_tokenizer_present and protein_tokenizer_index is not None:
    self.attributes.insert(protein_tokenizer_index, "protein_tokenizer")
```

Hope this would help to make a more stable vision LLMs!

And I'll rebase the main branch and see how it going. Thank you again!

@zhoubay yeah, but we don't expect users to change the order of attributes unless they are doing some advanced experiments and adding new attributes. In that case we'll expect them to also change the custom saving and loading. Usually we just have a tokenizer + image processor, one for each attribute, so not much opportunity to experiment 😄 

For Evolla feel free to do what you think works, I am not familiar with the model. You will get a review from Cyril on that later

@zucchini-nlp Thank you for your information and your collaboration! I'm looking forward to it!

Hey @zhoubay! Yes, we do enforce the use of modular now, especially as it is very much applicable in your case (a lot of code seem directly copy-pasted from other models)! In the end, it is a huge win-win scenario, as it greatly facilitates our review and maintenance process, and in turn ensures your model will get all potential future features/improvements we may roll out (which is otherwise not the case for all models, as model are supposed to stay more or less static in time) 🤗

Also, not sure what you did but you now have more than 150 changed files! 🙃 You should revert to previous state!

Hi @Cyrilvallez ! Thank you for your information! I'll convert my codes using modules! 

As for the file changes, I think most of them are involved because I just rebased my branch to the current main branch. Could you show me some clues about how to change it back?

Hi @Cyrilvallez,  I hope this would help! Sorry I messed up the history of commit! I'll never never never do rebase blindly again 😭 

No worries, we've all been there 😆 Let me know when you integrated the modular files and are ready for a review!

Sure! Thank you! I'll let you know when I finish integrating the modular files!

Just a heads-up, if you make 2 modular it will likely not work right now, but https://github.com/huggingface/transformers/pull/36287 will fix it! I will merge asap

Hey @Cyrilvallez ! I've got a little question about modular! Since the modular relies on the codes from other people, what happens if they have changed the file making the performance differently? 

For example, in my case, I inherited the `llama` code long long ago, like `transformers:v4.30.2`. And at that time, the signiture in `LlamaDecoderLayer.forward` has no `position_embeddings`. see [here](https://github.com/huggingface/transformers/blob/66fd3a8d626a32989f4569260db32785c6cbf42a/src/transformers/models/llama/modeling_llama.py#L264-L272).

```python
def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[Tuple[torch.Tensor]] = None,
    output_attentions: Optional[bool] = False,
    use_cache: Optional[bool] = False,
) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
```

And at that time, the `decoder_layer` won't accept `position_embeddings` as input (see [here](https://github.com/huggingface/transformers/blob/66fd3a8d626a32989f4569260db32785c6cbf42a/src/transformers/models/llama/modeling_llama.py#L578-L585)), because they compute `position_embeddings` inside `LlamaAttention` (see [here](https://github.com/huggingface/transformers/blob/66fd3a8d626a32989f4569260db32785c6cbf42a/src/transformers/models/llama/modeling_llama.py#L201-L202)).

```python
  layer_outputs = decoder_layer(
      hidden_states,
      attention_mask=attention_mask,
      position_ids=position_ids,
      past_key_value=past_key_value,
      output_attentions=output_attentions,
      use_cache=use_cache,
  )
```

However, for now, the `decoder_layer` must accept the `position_embeddings` as input (see [here](https://github.com/huggingface/transformers/blob/31bb662db19355e56b3125ab5c48b780ec13111c/src/transformers/models/llama/modeling_llama.py#L594-L604)). 

```python
  layer_outputs = decoder_layer(
      hidden_states,
      attention_mask=causal_mask,
      position_ids=position_ids,
      past_key_value=past_key_values,
      output_attentions=output_attentions,
      use_cache=use_cache,
      cache_position=cache_position,
      position_embeddings=position_embeddings,
      **flash_attn_kwargs,
  )
```

My point is, for my case, I just need to implement the `EvollaModel.forward` in fine-grain, so I have to copy all of the codes of `LlamaModel.forward` at that time of point, and I won't change anything in `LlamaDecoderLayer`, so I can inherit it. 

So my code of running `decoder_layer` would like previous one. This is implemented by me, so when they (llama team) change detailed implementations, these lines wouldn't change.

```python
  layer_outputs = decoder_layer(
      hidden_states,
      attention_mask=attention_mask,
      position_ids=position_ids,
      past_key_value=past_key_value,
      output_attentions=output_attentions,
      use_cache=use_cache,
  )
```

But over some time, the llama team changed the signiture of `decoder_layer` and `position_embeddings` are compulsory (see [here](https://github.com/huggingface/transformers/blob/31bb662db19355e56b3125ab5c48b780ec13111c/src/transformers/models/llama/modeling_llama.py#L319-L330)), my codes would suffer because of inheriting their `LlamaDecoderLayer` and wouldn't work. 

```python
def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[Cache] = None,
    output_attentions: Optional[bool] = False,
    use_cache: Optional[bool] = False,
    cache_position: Optional[torch.LongTensor] = None,
    position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
    **kwargs: Unpack[FlashAttentionKwargs],
) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
```

I don't know if I made myself clear, if you have any question about what I said, please let me know!

Looking forward to your idea to solve this issue! Thank you in advance!

Hi @Cyrilvallez , I made conversion to `modular` for `llama` parts and I found there's a typo in the doc, as issue #36362 pointed out. 

I fixed the typo in the doc, hope this wouldn't bother you!

Let me know if I should do more things!

Hi @Cyrilvallez ! I hope this message finds you well! It's been a week since my last commit, and I wanted to kindly remind you to take a look at the PR when you have a moment. Please let me know if there are any issues or concerns you might have. Looking forward to your feedback!

Thank you Cyrilvallez for your suggestions! I'm working on your comments!

Hi @Cyrilvallez , I've been distracted by others things for a while! I think I've finished all things you've mentioned, but something tests (not evolla) still sadly failed. Would you take some time take a look at my Evolla and check if there's still something going wrong? Thank you very much for your time!

Hi @Cyrilvallez ! I think evolla has passed all tests, could you please have some time check on it to see if we can merge into the main branch or not? Thank you for your time!

Hi @Cyrilvallez and @ArthurZucker ! Thank you for your excellent work! I've fixed the issue Arthur pointed out! Please let me know what else should I do!

Hi @Cyrilvallez and @ArthurZucker , sorry to bother you in your busy time! It's been over two weeks since last review, could you please spend some extra time on this PR? I would be more than grateful for your time! Thank you!

Hi @Cyrilvallez @ArthurZucker , it's been a month since last review :), would you be so kind to take some time to look at this PR and give me some information that I could know which direction I should head for! Thank you!

Also, I just updated the branch because we changed some rules in our linter etc upstream, just running `make style` should fix the quality CI that will be red!

Hi @Cyrilvallez , thank you for your hard work! I've finished all your suggestions except the two:
1. Renaming `RotaryEmbedding` to `EvollaSaProtRotaryEmbedding`
2. Changing the implementation way of mask in `EvollaSaProtProteinEncoder`

Please take some time to revise the according comments and give me further suggesions 😄 

Thank you again for your hard work!

Also, you have unrelated pickle files to remove

Hi @Cyrilvallez ! Thank you for your hard work! I've finished all of your suggestions! Hope to see this to merge! 

And for the errors in CI, I think those are not caused by Evolla. Hope this doesn't matter!

Let me know if I have further things to do!

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

Alright @zhoubay, I pushed a few commits myself to try to cleanup as much as possible! Can you check if this still works correctly with your checkpoints? Also, you need to fill the docstrings in modular for the main models 🤗

run-slow: evolla

This comment contains run-slow, running the specified jobs: 

models: ['models/evolla']
quantizations: [] ...

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: auto, evolla

Hi @Cyrilvallez ! Thank you for your hard work! I've made some commit to adjust, and checked the performance of Evolla, it's as good as before!