[WIP] Implements Roberta Model #679
Conversation
| distance_embedding = None | ||
| position_embedding_type = config.position_embedding_type | ||
|
|
||
| if position_embedding_type == "relative_key" or position_embedding_type == "relative_key_query": |
| key = None | ||
| ) -> Tuple[NamedArray]: | ||
|
|
||
| query_layer = self.transpose_for_scores(self.q_proj(hidden_states)) |
There was a problem hiding this comment.
in theory you shouldn't need this transpose_for_scores
There was a problem hiding this comment.
I did it for compatibility with huggingface. I noticed that the llama code has the linear layer output it directly in the correct shape, but I didnt want to deal with communicating that in the state_dict functions. Do you want me to change it?
There was a problem hiding this comment.
parameters are best declared in the same order as they are in HF (though that can be worked around) but intermediate values like query_layer can be in any order really. Haliax will automatically transpose things as needed.
|
|
||
| attention_scores /= jnp.sqrt(self.HeadSize.size) | ||
|
|
||
| if attention_mask is not None: |
There was a problem hiding this comment.
masks in Levanter are traditionally binary, which means you need to use something like hax.where(attention_mask, attention_scores, -1E9)
src/levanter/models/testing.ipynb
Outdated
There was a problem hiding this comment.
Will do. Do you want me to add it somewhere else in levanter or just keep it to myself?
There was a problem hiding this comment.
you can just move it into examples/roberta/ or something. Ideally it was be proper unit tests
| return q_embed, k_embed | ||
|
|
||
|
|
||
| def llama_rotary_pos_emb( |
There was a problem hiding this comment.
you'll want to delete this i think? (if you end up needing it, just import from llama.py)
|
I should add: looking good! I know this ended up being a big lift and I appreciate you all taking it on! |
RobertaForMaskedLM
…_hidden_states implementation in jax model
… into roberta-model
…ention mask more robust
No description provided.