Delete files world_model/modeling_llama_action.py world_model/configuration_llama_action.py with huggingface_hub

world_model/configuration_llama_action.py

@@ -1,13 +0,0 @@
-from transformers import LlamaConfig
-
-
-class LlamaActionConfig(LlamaConfig):
-    model_type = "llama_action"
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.num_spatio_embeddings = kwargs.get("num_spatio_embeddings", 582)
-        self.num_temporal_embeddings = kwargs.get("num_temporal_embeddings", 25)
-        self.num_action_embeddings = kwargs.get("num_action_tokens", 5)
-        self.num_image_patches = kwargs.get("num_image_patches", 576)
-        self.action_dim = kwargs.get("action_dim", 3)

world_model/modeling_llama_action.py

@@ -1,237 +0,0 @@
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-from transformers import LlamaForCausalLM
-from transformers.modeling_outputs import CausalLMOutputWithPast
-
-from .configuration_llama_action import LlamaActionConfig
-
-
-class LearnableFactorizedSpatioTemporalPositionalEmbedding(nn.Module):
-    def __init__(self, num_spatio_embeddings: int, num_temporal_embeddings: int, embedding_dim: int):
-        super().__init__()
-        self.spatio_embeddings = nn.Embedding(num_spatio_embeddings, embedding_dim)
-        self.temporal_embeddings = nn.Embedding(num_temporal_embeddings, embedding_dim)
-        self.num_spatio_embeddings = num_spatio_embeddings
-        self.num_temporal_embeddings = num_temporal_embeddings
-
-    def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int):
-        seq_length = attention_mask.size(1)
-        batch_size = attention_mask.size(0)
-
-        if past_key_values_length == 0:
-            # create a tensor of the form [0, 1, 2, ..., num_spatio_embeddings-1]
-            spatio_indices = torch.arange(
-                self.num_spatio_embeddings,
-                device=attention_mask.device
-            ).repeat(self.num_temporal_embeddings).unsqueeze(0).repeat((batch_size, 1))
-
-            # create a tensor of the form [0, 0, 0, ..., 1, 1, 1, ..., 2, 2, 2, ...]
-            temporal_indices = torch.arange(
-                self.num_temporal_embeddings,
-                device=attention_mask.device
-            ).repeat_interleave(self.num_spatio_embeddings).unsqueeze(0).repeat((batch_size, 1))
-
-            spatio_indices = spatio_indices[:, :seq_length]
-            temporal_indices = temporal_indices[:, :seq_length]
-            
-        else:
-            temporal_index = past_key_values_length // self.num_spatio_embeddings
-            spatio_index = past_key_values_length % self.num_spatio_embeddings
-            spatio_indices = torch.tensor([[spatio_index]], device=attention_mask.device).repeat((batch_size, 1))
-            temporal_indices = torch.tensor([[temporal_index]], device=attention_mask.device).repeat((batch_size, 1))
-
-        return self.spatio_embeddings(spatio_indices) + self.temporal_embeddings(temporal_indices)
-
-
-class LlamaActionForCausalLM(LlamaForCausalLM):
-    config_class = LlamaActionConfig
-
-    def __init__(self, config: LlamaActionConfig):
-        super().__init__(config)
-
-        self.num_spatio_embeddings = config.num_spatio_embeddings
-        self.num_temporal_embeddings = config.num_temporal_embeddings
-        self.num_image_patches = config.num_image_patches
-        self.num_action_embeddings = config.num_action_embeddings
-
-        self.pos_embedding_spatio_temporal = LearnableFactorizedSpatioTemporalPositionalEmbedding(
-            config.num_spatio_embeddings, config.num_temporal_embeddings, config.hidden_size,
-        )
-
-        self.action_projection = nn.Linear(config.action_dim, config.hidden_size)
-
-        self.post_init()
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        actions: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        past_key_values: Optional[List[torch.Tensor]] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithPast]:
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        if labels is not None:
-            use_cache = False
-
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError(
-                "You cannot specify both input_ids and inputs_embeds at the same time"
-            )
-        elif input_ids is not None:
-            pass
-        elif inputs_embeds is not None:
-            pass
-        else:
-            raise ValueError("You have to specify either input_ids or inputs_embeds")
-
-        inputs_embeds = self.model.get_input_embeddings()(input_ids)
-        if past_key_values is None or len(past_key_values) == 0:
-            inputs_embeds_list = torch.split(
-                inputs_embeds,
-                split_size_or_sections=self.num_image_patches,
-                dim=1
-            )
-            actions_list = torch.split(
-                actions,
-                split_size_or_sections=self.num_action_embeddings,
-                dim=1
-            )
-
-            embeddings = []
-            if len(inputs_embeds_list) == len(actions_list):
-                # mostly used in training phase
-                for inputs_embeds, action_embeds in zip(inputs_embeds_list, actions_list):
-                    action_features = self.action_projection(action_embeds)
-                    embeddings.append(inputs_embeds)
-                    embeddings.append(action_features)
-            elif len(inputs_embeds_list) < len(actions_list):
-                # used in inference phase (mostly)
-                for i, inputs_embeds in enumerate(inputs_embeds_list):
-                    embeddings.append(inputs_embeds)
-                    if i < len(inputs_embeds_list) - 1:
-                        # the last frame might be generating image tokens, so we don't add action embedding
-                        action_embeds = self.action_projection(actions_list[i])
-                        embeddings.append(action_embeds)
-                if inputs_embeds_list[-1].size(1) == self.num_image_patches:
-                    # if the last frame has generated all image tokens, we add action embedding
-                    action_embeds = self.action_projection(actions_list[len(inputs_embeds_list) - 1])
-                    embeddings.append(action_embeds)
-        else:
-            if isinstance(past_key_values, tuple):
-                past_key_values_length = past_key_values[0][0].size(2)
-            else:
-                past_key_values_length = past_key_values.get_seq_length()
-            embeddings = []
-            # create an interleaved sequence of image and action embeddings like image, image, ..., image, action, action, ..., action
-            # we only generate image tokens, so we add action tokens after generating one frame
-            if past_key_values_length % self.num_spatio_embeddings == (self.num_spatio_embeddings - self.num_action_embeddings):
-                seq_index = past_key_values_length // self.num_spatio_embeddings + 1
-                actions_list = torch.split(
-                    actions,
-                    split_size_or_sections=self.num_action_embeddings,
-                    dim=1
-                )
-                action_features = self.action_projection(actions_list[seq_index - 1])
-                embeddings.append(action_features)
-                embeddings.append(inputs_embeds)
-            else:
-                pass
-
-        if len(embeddings) > 0:
-            inputs_embeds = torch.cat(embeddings, dim=1)
-
-        # insert spatio-temporal positional embedding
-        if past_key_values is not None:
-            if isinstance(past_key_values, tuple):
-                past_key_values_length = past_key_values[0][0].size(2)
-            else:
-                past_key_values_length = past_key_values.get_seq_length()
-        else:
-            past_key_values_length = 0
-        inputs_embeds += self.pos_embedding_spatio_temporal(inputs_embeds, past_key_values_length)
-
-        outputs = self.model(
-            input_ids=None,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        sequence_output = outputs[0]
-        logits = self.lm_head(sequence_output).contiguous()
-
-        loss = None
-        if labels is not None:
-            shift_logits = logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            loss_fct = nn.CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
-
-        if not return_dict:
-            output = (logits,) + outputs[1:]
-            return ((loss,) + output) if loss is not None else output
-
-        return CausalLMOutputWithPast(
-            loss=loss,
-            logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-    def prepare_inputs_for_generation(
-        self,
-        input_ids,
-        past_key_values=None,
-        attention_mask=None,
-        use_cache=None,
-        **kwargs):
-        batch_size = input_ids.size(0)
-        seq_length = input_ids.size(1)
-        n_frames = seq_length // self.num_image_patches
-        attention_mask_length = n_frames * (self.num_image_patches + self.num_action_embeddings)
-        if seq_length % self.num_image_patches != 0:
-            n_last_frame_tokens = seq_length % self.num_image_patches
-            attention_mask_length += n_last_frame_tokens
-        else:
-            print(f"attempting to generate new frame - frame no: {n_frames + 1}")
-        attention_mask = torch.ones((batch_size, attention_mask_length), device=input_ids.device, dtype=torch.long)
-        # cut decoder_input_ids if past_key_values is used
-        if past_key_values is not None and len(past_key_values) > 0:
-            if isinstance(past_key_values, tuple):
-                past_length = past_key_values[0][0].size(2)
-            else:
-                past_length = past_key_values.get_seq_length()
-            if input_ids.size(1) > past_length:
-                remove_prefix_length = past_length
-            else:
-                remove_prefix_length = input_ids.size(1) - 1
-            input_ids = input_ids[:, remove_prefix_length:]
-            seq_length = input_ids.size(1)
-            past_key_values_length = past_length
-            mask_seq_length = seq_length + past_key_values_length
-            if past_key_values_length % self.num_spatio_embeddings == (self.num_spatio_embeddings - self.num_action_embeddings):
-                mask_seq_length += self.num_action_embeddings
-            attention_mask = torch.ones((batch_size, mask_seq_length), device=input_ids.device, dtype=torch.long)
-
-        return {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "actions": kwargs.get("actions"),
-            "past_key_values": past_key_values,
-            "use_cache": use_cache,
-        }