Upload modeling_olmo.py with huggingface_hub

modeling_olmo.py

@@ -0,0 +1,187 @@
+from dataclasses import fields
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import math
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
+from transformers.models.auto import AutoModelForCausalLM
+
+from .config import ModelConfig
+from .model import OLMo
+
+from .configuration_olmo import OLMoConfig
+
+def create_model_config_from_pretrained_config(config: OLMoConfig):
+    """
+    Utility function
+    """
+
+    kwargs = {}
+    for field in fields(ModelConfig):
+        kwargs[field.name] = getattr(config, field.name)
+
+    model_config = ModelConfig(**kwargs)
+    return model_config
+
+class OLMoPreTrainedModel(PreTrainedModel):
+    config_class = OLMoConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["OLMoBlock"]
+    # _skip_keys_device_placement = ["past_key_values", "causal_mask"]
+    _skip_keys_device_placement = ["past_key_values"]
+
+    def _init_weights(self, module):
+        # `OLMoModel.reset_parameters` initializes weights of itself and its children
+        if isinstance(module, OLMo):
+            module.reset_parameters()
+
+class OLMoForCausalLM(OLMoPreTrainedModel):
+    _tied_weights_keys = []
+    # _tied_weights_keys = ["transformer.wte.weight"]
+
+    def __init__(self, config: OLMoConfig):
+        super().__init__(config)
+        self.model = OLMo(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.transformer.wte
+
+    def set_input_embeddings(self, value: torch.nn.Module):
+        self.model.transformer.wte = value
+
+    def get_output_embeddings(self):
+        if self.config.weight_tying:
+            return self.model.transformer.wte
+        else:
+            return self.model.transformer.ff_out
+
+    def set_output_embeddings(self, value: torch.nn.Module):
+        if self.config.weight_tying:
+            self.model.transformer.wte = value
+        else:
+            self.model.transformer.ff_out = value
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        attention_bias: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        Returns:
+        Example:
+        ```python
+        >>> from transformers import AutoTokenizer, OLMoForCausalLM
+        >>> model = OLMoForCausalLM.from_pretrained("allenai/OLMo-7B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-7B")
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        output_attentions = output_attentions or self.config.output_attentions
+        output_hidden_states = output_hidden_states or self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        assert not output_attentions
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        base_output: Union[BaseModelOutputWithPast, Tuple] = self.model.forward(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            attention_bias=attention_bias,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_hidden_states=output_hidden_states,
+        )
+
+        last_hidden_state = base_output.last_hidden_state if return_dict else base_output[0]
+
+        # Get logits.
+        # shape: (batch_size, seq_len or 1, vocab_size)
+        if self.config.weight_tying:
+            logits = F.linear(last_hidden_state, self.model.transformer.wte.weight, None)  # type: ignore
+        else:
+            logits = self.model.transformer.ff_out(last_hidden_state)  # type: ignore
+        if self.config.scale_logits:
+            logits.mul_(1 / math.sqrt(self.config.d_model))
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = torch.nn.CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.embedding_size) # changed to self.config.embedding_size from self.config.vocab_size
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + base_output[1:]
+            return (loss,) + output if loss is not None else output
+
+        assert isinstance(base_output, BaseModelOutputWithPast)
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=base_output.past_key_values,
+            hidden_states=base_output.hidden_states,
+            attentions=base_output.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple]] = None, **kwargs
+    ):
+        if past_key_values:
+            # This is because we want the model to only process the last generated token.
+            input_ids = input_ids[:, -1:]
+        model_inputs = {"input_ids": input_ids, "past_key_values": past_key_values}
+
+        if 'cache_position' in kwargs: kwargs.pop("cache_position")
+        if past_key_values and ("input_embeds" in kwargs or "inputs_embeds" in kwargs): kwargs.pop("inputs_embeds")
+        model_inputs.update(kwargs)
+        # logger.warning("%s %s", kwargs.keys(), model_inputs.keys())
+        # model_inputs["use_cache"] = kwargs.pop("use_cache", self.config.use_cache)
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
+
+# Register the model so that it is available for transformer pipelines, auto-loading, etc.
+# AutoModelForCausalLM.register(OLMoConfig, OLMoForCausalLM)