Delete CHATTS/model/gpt.py

CHATTS/model/gpt.py

@@ -1,265 +0,0 @@
-import os
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-import logging
-from tqdm import tqdm
-from einops import rearrange
-from transformers.cache_utils import Cache
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.nn.utils.parametrize as P
-from torch.nn.utils.parametrizations import weight_norm
-from transformers import LlamaModel, LlamaConfig
-    
-    
-class LlamaMLP(nn.Module):
-    def __init__(self, hidden_size, intermediate_size):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = F.silu
-
-    def forward(self, x):
-        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
-    
-    
-class GPT_warpper(nn.Module):
-    def __init__(
-        self, 
-        gpt_config, 
-        num_audio_tokens,
-        num_text_tokens,
-        num_vq=4,
-        **kwargs,
-        ):
-        super().__init__()
-
-        self.logger = logging.getLogger(__name__)
-        self.gpt = self.build_model(gpt_config)
-        self.model_dim = self.gpt.config.hidden_size 
-
-        self.num_vq = num_vq
-        self.emb_code = nn.ModuleList([nn.Embedding(num_audio_tokens, self.model_dim) for i in range(self.num_vq)])
-        self.emb_text = nn.Embedding(num_text_tokens, self.model_dim)
-        self.head_text = weight_norm(nn.Linear(self.model_dim, num_text_tokens, bias=False), name='weight')
-        self.head_code = nn.ModuleList([weight_norm(nn.Linear(self.model_dim, num_audio_tokens, bias=False), name='weight') for i in range(self.num_vq)])
-
-    def build_model(self, config):
-        
-        configuration = LlamaConfig(**config)
-        model = LlamaModel(configuration)
-        del model.embed_tokens
-        
-        return model
-    
-    def get_emb(self, input_ids, text_mask, **kwargs):
-
-        emb_text = self.emb_text(input_ids[text_mask][:, 0])
-        
-        emb_code = [self.emb_code[i](input_ids[~text_mask][:, i]) for i in range(self.num_vq)]
-        emb_code = torch.stack(emb_code, 2).sum(2)
-        
-        emb = torch.zeros((input_ids.shape[:-1])+(emb_text.shape[-1],), device=emb_text.device, dtype=emb_text.dtype)
-        emb[text_mask] = emb_text
-        emb[~text_mask] = emb_code.to(emb.dtype)
-        
-        return emb
-    
-    def prepare_inputs_for_generation(
-        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs
-    ):
-        # With static cache, the `past_key_values` is None
-        # TODO joao: standardize interface for the different Cache classes and remove of this if
-        has_static_cache = False
-        if past_key_values is None:
-            past_key_values = getattr(self.gpt.layers[0].self_attn, "past_key_value", None)
-            has_static_cache = past_key_values is not None
-
-        past_length = 0
-        if past_key_values is not None:
-            if isinstance(past_key_values, Cache):
-                past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
-                max_cache_length = (
-                    torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
-                    if past_key_values.get_max_length() is not None
-                    else None
-                )
-                cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
-            # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
-            else:
-                cache_length = past_length = past_key_values[0][0].shape[2]
-                max_cache_length = None
-
-            # Keep only the unprocessed tokens:
-            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
-            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
-            # input)
-            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
-                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
-            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
-            # input_ids based on the past_length.
-            elif past_length < input_ids.shape[1]:
-                input_ids = input_ids[:, past_length:]
-            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
-
-            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
-            if (
-                max_cache_length is not None
-                and attention_mask is not None
-                and cache_length + input_ids.shape[1] > max_cache_length
-            ):
-                attention_mask = attention_mask[:, -max_cache_length:]
-
-        position_ids = kwargs.get("position_ids", None)
-        if attention_mask is not None and position_ids is None:
-            # create position_ids on the fly for batch generation
-            position_ids = attention_mask.long().cumsum(-1) - 1
-            position_ids.masked_fill_(attention_mask == 0, 1)
-            if past_key_values:
-                position_ids = position_ids[:, -input_ids.shape[1] :]
-
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and past_key_values is None:
-            model_inputs = {"inputs_embeds": inputs_embeds}
-        else:
-            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
-            # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
-            # TODO: use `next_tokens` directly instead.
-            model_inputs = {"input_ids": input_ids.contiguous()}
-
-        input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
-        if cache_position is None:
-            cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
-        else:
-            cache_position = cache_position[-input_length:]
-
-        if has_static_cache:
-            past_key_values = None
-
-        model_inputs.update(
-            {
-                "position_ids": position_ids,
-                "cache_position": cache_position,
-                "past_key_values": past_key_values,
-                "use_cache": kwargs.get("use_cache"),
-                "attention_mask": attention_mask,
-            }
-        )
-        return model_inputs
-    
-    def generate(
-        self, 
-        emb, 
-        inputs_ids, 
-        temperature, 
-        eos_token, 
-        attention_mask = None,
-        max_new_token = 2048, 
-        min_new_token = 0,
-        LogitsWarpers = [],
-        LogitsProcessors = [],
-        infer_text=False,
-        return_attn=False,
-        return_hidden=False,
-    ):
-        
-        with torch.no_grad():   
-        
-            attentions = []
-            hiddens = []
-            
-            start_idx, end_idx = inputs_ids.shape[1], torch.zeros(inputs_ids.shape[0], device=inputs_ids.device, dtype=torch.long)
-            finish = torch.zeros(inputs_ids.shape[0], device=inputs_ids.device).bool()
-            
-            temperature = temperature[None].expand(inputs_ids.shape[0], -1)
-            temperature = rearrange(temperature, "b n -> (b n) 1")
-
-            attention_mask_cache = torch.ones((inputs_ids.shape[0], inputs_ids.shape[1]+max_new_token,), dtype=torch.bool, device=inputs_ids.device)
-            if attention_mask is not None:
-                attention_mask_cache[:, :attention_mask.shape[1]] = attention_mask
-            
-            for i in tqdm(range(max_new_token)):
-        
-                model_input = self.prepare_inputs_for_generation(inputs_ids, 
-                    outputs.past_key_values if i!=0 else None, 
-                    attention_mask_cache[:, :inputs_ids.shape[1]], use_cache=True)
-            
-                if i == 0:
-                    model_input['inputs_embeds'] = emb
-                else:
-                    if infer_text:
-                        model_input['inputs_embeds'] = self.emb_text(model_input['input_ids'][:,:,0])
-                    else:
-                        code_emb = [self.emb_code[i](model_input['input_ids'][:,:,i]) for i in range(self.num_vq)]
-                        model_input['inputs_embeds'] = torch.stack(code_emb, 3).sum(3)
-                
-                model_input['input_ids'] = None
-                outputs = self.gpt.forward(**model_input, output_attentions=return_attn)
-                attentions.append(outputs.attentions)
-                hidden_states = outputs[0] # 🐻
-                if return_hidden:
-                    hiddens.append(hidden_states[:, -1])
-
-                with P.cached():
-                    if infer_text:
-                        logits = self.head_text(hidden_states) 
-                    else:
-                        logits = torch.stack([self.head_code[i](hidden_states) for i in range(self.num_vq)], 3)
-        
-                logits = logits[:, -1].float()
-
-                if not infer_text:
-                    logits = rearrange(logits, "b c n -> (b n) c")
-                    logits_token = rearrange(inputs_ids[:, start_idx:], "b c n -> (b n) c")
-                else:
-                    logits_token = inputs_ids[:, start_idx:, 0]
-                    
-                logits = logits / temperature
-                
-                for logitsProcessors in LogitsProcessors:
-                    logits = logitsProcessors(logits_token, logits)
-                    
-                for logitsWarpers in LogitsWarpers:
-                    logits = logitsWarpers(logits_token, logits)
-                    
-                if i < min_new_token:
-                    logits[:, eos_token] = -torch.inf
-                
-                scores = F.softmax(logits, dim=-1)
-            
-                idx_next = torch.multinomial(scores, num_samples=1)
-                
-                if not infer_text:
-                    idx_next = rearrange(idx_next, "(b n) 1 -> b n", n=self.num_vq)
-                    finish = finish | (idx_next == eos_token).any(1)
-                    inputs_ids = torch.cat([inputs_ids, idx_next.unsqueeze(1)], 1)
-                else:
-                    finish = finish | (idx_next == eos_token).any(1)
-                    inputs_ids = torch.cat([inputs_ids, idx_next.unsqueeze(-1).expand(-1, -1, self.num_vq)], 1)
-
-                end_idx = end_idx + (~finish).int()
-            
-                if finish.all():
-                    break
-            
-            inputs_ids = [inputs_ids[idx, start_idx: start_idx+i] for idx, i in enumerate(end_idx.int())]
-            inputs_ids = [i[:, 0] for i in inputs_ids] if infer_text else inputs_ids
-            
-            if return_hidden:
-                hiddens = torch.stack(hiddens, 1)
-                hiddens = [hiddens[idx, :i] for idx, i in enumerate(end_idx.int())]
-                    
-            if not finish.all():
-                self.logger.warn(f'Incomplete result. hit max_new_token: {max_new_token}')    
-                   
-            return {
-                'ids': inputs_ids, 
-                'attentions': attentions,
-                'hiddens':hiddens,
-            }