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.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} kwargs.pop("cache_position") 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)