llama_model_v1.py

import math
from dataclasses import dataclass
import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F
from typing_extensions import Self
from typing import Optional
from transformers.modeling_utils import PreTrainedModel
from torch.distributions import Categorical
import torch.nn.functional as F


@dataclass
class LLaMAHFConfig:
    block_size: int = 78
    n_layer: int = 32
    n_head: int = 32
    n_embd: int = 4096
    T5_xxl_dim: int = 768

    @classmethod
    def from_name(cls, name: str) -> Self:
        return cls(**llama_configs[name])


llama_configs = {
    "Normal_size": dict(n_layer=12, n_head=12, n_embd=768)
}


class LLaMAHF(nn.Module):
    def __init__(self, config: LLaMAHFConfig, num_diffusion_head_layers=9, input_token_dim=16, device=torch.device('cuda'), width=1792) -> None:
        super().__init__()
        assert config.block_size is not None
        self.config = config

        cond_dim = config.T5_xxl_dim
         
        self.transformer = nn.ModuleDict(
            dict(
                wte=nn.Linear(input_token_dim, config.n_embd),   
                cond_embed=nn.Linear(cond_dim, config.n_embd),  
                h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
                ln_f=RMSNorm(config.n_embd),
                )
            )
        
        target_channels = input_token_dim
        from models.diffloss import DiffLoss
        self.diff_loss = DiffLoss(
                target_channels=target_channels,
                z_channels=config.n_embd,
                width=width,
                depth=num_diffusion_head_layers,
                num_sampling_steps='50',
                grad_checkpointing=False,
            )
        self.diff_loss = self.diff_loss.to(device)
        self.out_proj = nn.Linear(config.n_embd, config.n_embd)
        self.use_out_proj = True
        

    def _tie_or_clone_weights(self, output_embeddings, input_embeddings):
        """Tie or clone module weights depending of whether we are using TorchScript or not"""
        output_embeddings.weight = input_embeddings.weight

        if getattr(output_embeddings, "bias", None) is not None:
            output_embeddings.bias.data = nn.functional.pad(
                output_embeddings.bias.data,
                (
                    0,
                    output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0],
                ),
                "constant",
                0,
            )
        if hasattr(output_embeddings, "out_features") and hasattr(input_embeddings, "num_embeddings"):
            output_embeddings.out_features = input_embeddings.num_embeddings

    def get_input_embeddings(self):
        return self.transformer.wte
    
    def set_input_embeddings(self, value):
        self.transformer.wte = value

    def get_output_embeddings(self):
        return self.lm_head
    
    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02 / math.sqrt(2 * self.config.n_layer))
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02 / math.sqrt(2 * self.config.n_layer))
            
    

    def forward_sample(self, idx: torch.Tensor, clip_feature: torch.Tensor, y_mask) -> torch.Tensor:
        
        text_length = clip_feature.shape[1]
        if len(idx) == 0:
            x = self.llama_proj(clip_feature)[:, :int(y_mask[0].sum()), :]
        else:
            _, t = idx.size()
            assert (
                t <= self.config.block_size
            ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
            # forward the LLaMA model itself
            x = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)
            x = torch.cat((self.llama_proj(clip_feature)[:, :int(y_mask[0].sum()), :],x), dim=1)

        for block in self.transformer.h:
            x = block(x, y_mask)
        x = self.transformer.ln_f(x)
        logits = x
        return logits

    

    def sample_for_eval_CFG(self, text, length=196, tokenize_model=None, device=torch.device('cuda'), unit_length=4, cfg=4.0):
        max_token_len = length // unit_length
        for k in range(max_token_len): 
            if k == 0:
                x = []
            else:
                x = xs
            
            feat_text = torch.from_numpy(tokenize_model.encode(text)).float()
            feat_text = feat_text.to(device)
            conditions = self.forward(x, feat_text)
            conditions = conditions[:, -1, :]                            
            
            empty_text = ''
            empty_feat_text = torch.from_numpy(tokenize_model.encode(empty_text)).float()   
            empty_feat_text = empty_feat_text.unsqueeze(0)
            empty_feat_text = empty_feat_text.to(device)
            empty_conditions = self.forward(x, empty_feat_text)       
            empty_conditions = empty_conditions[:, -1, :]                             
            temperature = 1.0
            
            # chunk
            if cfg != 1:
                mix_conditions = torch.cat([conditions, empty_conditions], dim=0)     
                sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg)  
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)          
            else:   # no cfg
                scaled_logits = self.diff_loss.sample(conditions, temperature=temperature, cfg=1) 

            scaled_logits = scaled_logits.unsqueeze(0)     
            
            if k == 0:
                xs = scaled_logits
            else:
                xs = torch.cat((xs, scaled_logits), dim=1)

        return xs
    
    
    
    # For inference, can stop sampling when the distance between the current token and the reference end token is less than the threshold.
    def sample_for_eval_CFG_inference(self, text, length=312, tokenizer=None, device=torch.device('cuda'), unit_length=4, reference_end_latent=None, threshold=0.1, cfg=4.0, temperature=1.0):
        max_token_len = length // unit_length
        feat_text = torch.from_numpy(tokenizer.encode(text)).float()
        feat_text = feat_text.to(device)

        # CFG inference
        empty_text = ''
        empty_feat_text = torch.from_numpy(tokenizer.encode(empty_text)).float()   # torch.Size([32, 768])
        empty_feat_text = empty_feat_text.unsqueeze(0)
        empty_feat_text = empty_feat_text.to(device)
        
        for k in range(max_token_len): 
            if k == 0:
                x = []
            else:
                x = xs
                
            conditions = self.forward_inference(x, feat_text)
            conditions = conditions[:, -1, :]                            

            empty_conditions = self.forward(x, empty_feat_text)       
            empty_conditions = empty_conditions[:, -1, :]                             
            
            mix_conditions = torch.cat([conditions, empty_conditions], dim=0)     
            sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg)  

            # chunk
            if cfg != 1:
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)           
            else:
                scaled_logits = sampled_token_latent

            scaled_logits = scaled_logits.unsqueeze(0)      
            
            if reference_end_latent is not None:
                distance_l2 = torch.sqrt(torch.sum((scaled_logits - reference_end_latent)**2))
                print(distance_l2)
                if distance_l2 < threshold:
                    break

            if k == 0:
                xs = scaled_logits
            else:
                xs = torch.cat((xs, scaled_logits), dim=1)

        return xs

        
    def sample_for_eval_CFG_inference2(self, feat_clip_text, empty_feat_clip_text, if_categorial=False, length=312, clip_model=None, device=torch.device('cuda'), tokenizer='clip', unit_length=4, reference_end_token=None, threshold=3, cfg=4.5, temperature=1.0):
       
        import clip
        max_token_len = length // unit_length
        
        for k in range(max_token_len): 
            if k == 0:
                x = []
            else:
                x = xs

            try:
                conditions = self.forward(x, feat_clip_text)
            except:
                conditions = self.forward(x, feat_clip_text.unsqueeze(0))
            
            
            conditions = conditions[:, -1, :]                            
            
            

            empty_conditions = self.forward(x, empty_feat_clip_text)       
            empty_conditions = empty_conditions[:, -1, :]                             
            
            mix_conditions = torch.cat([conditions, empty_conditions], dim=0)     
            sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg)  

            # chunk
            if cfg != 1:
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)           
            else:
                scaled_logits = sampled_token_latent

            scaled_logits = scaled_logits.unsqueeze(0)      
            
            if reference_end_token is not None:
                distance_l2 = torch.sqrt(torch.sum((scaled_logits - reference_end_token)**2))
                print(distance_l2)
                if distance_l2 < threshold:
                    break

            if k == 0:
                xs = scaled_logits
            else:
                xs = torch.cat((xs, scaled_logits), dim=1)

        return xs

    def sample_for_eval_CFG_inference_next_one(self, current_token=[], feat_clip_text=None, empty_feat_clip_text=None, if_categorial=False, length=312, clip_model=None, device=torch.device('cuda'), tokenizer='clip', unit_length=4, reference_end_token=None, threshold=3, cfg=4.5, temperature=1.0):
        
        import clip
        max_token_len = length // unit_length
       
        
        for k in range(1): 
            
            if current_token == []:
                x = []
            else:
                x = torch.cat(current_token, dim=1)
            
            
            try:
                conditions = self.forward(x, feat_clip_text)
            except:
                conditions = self.forward(x, feat_clip_text.unsqueeze(0))
            
            
            conditions = conditions[:, -1, :]                             
            

            empty_conditions = self.forward(x, empty_feat_clip_text)      
            empty_conditions = empty_conditions[:, -1, :]                         

            mix_conditions = torch.cat([conditions, empty_conditions], dim=0)   
            sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg)  

            # chunk
            if cfg != 1:
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)         
            else:
                scaled_logits = sampled_token_latent


            scaled_logits = scaled_logits.unsqueeze(0)    
            

            if k == 0:
                xs = scaled_logits
            else:
                xs = torch.cat((xs, scaled_logits), dim=1)

        return xs
    

    def sample_for_eval_CFG_babel(self, A_text, B_text, A_motion, if_categorial=False, length=6400, clip_model=None, device=torch.device('cuda'), tokenizer='clip', unit_length=4, reference_end_token=None, cfg=7.0, threshold=3):

        import clip
        B_token_length = length // unit_length - A_motion.shape[0]

        if tokenizer == 'clip':
            A_text = clip.tokenize(A_text, truncate=True).to(device) 
            A_feat_clip_text = clip_model.encode_text(A_text).float()  
            B_text = clip.tokenize(B_text, truncate=True).to(device)
            B_feat_clip_text = clip_model.encode_text(B_text).float()
        elif tokenizer == 't5-xxl':
            A_feat_clip_text = torch.from_numpy(clip_model.encode(A_text)).float()
            A_feat_clip_text = A_feat_clip_text.to(device)
            B_feat_clip_text = torch.from_numpy(clip_model.encode(B_text)).float()
            B_feat_clip_text = B_feat_clip_text.to(device)
        
        A_text_embeddings = self.transformer.cond_embed(A_feat_clip_text).unsqueeze(0)
        B_text_embeddings = self.transformer.cond_embed(B_feat_clip_text).unsqueeze(0)

        A_motion = A_motion.unsqueeze(0)
        A_motion_embeddings = self.transformer.wte(A_motion)
        B_motion = torch.tensor([]).to(device)

        for k in range(B_token_length): 
            if k == 0:
                x = torch.cat([A_text_embeddings, A_motion_embeddings, B_text_embeddings], dim=1)
            else:
                x = xs

            
            conditions = self.forward_babel_eval(x)
            conditions = conditions[:, -1, :]                       

            empty_clip_text = ''
            if tokenizer == 'clip':
                empty_text = clip.tokenize(empty_clip_text, truncate=True).to(device)
                empty_feat_clip_text = clip_model.encode_text(empty_text).float()
            elif tokenizer == 't5-xxl':
                empty_feat_clip_text = torch.from_numpy(clip_model.encode(empty_clip_text)).float()   
                empty_feat_clip_text = empty_feat_clip_text.unsqueeze(0)
                empty_feat_clip_text = empty_feat_clip_text.to(device)

            empty_feat_clip_text_embedding = self.transformer.cond_embed(empty_feat_clip_text).unsqueeze(0)
            
            if k == 0:
                empty_input = torch.cat([empty_feat_clip_text_embedding, A_motion_embeddings, empty_feat_clip_text_embedding], dim=1)
                empty_conditions = self.forward_babel_eval(empty_input)
            else:
                B_motion_embeddings = self.transformer.wte(B_motion)
                empty_input = torch.cat([empty_feat_clip_text_embedding, A_motion_embeddings, empty_feat_clip_text_embedding, B_motion_embeddings], dim=1)
                empty_conditions = self.forward_babel_eval(empty_input)      
            
            empty_conditions = empty_conditions[:, -1, :]                            
            temperature = 1.0
            
            mix_conditions = torch.cat([conditions, empty_conditions], dim=0)    
            sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg)  

            # chunk
            if cfg != 1:
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)           
            else:
                scaled_logits = sampled_token_latent


            scaled_logits = scaled_logits.unsqueeze(0)      
            

            B_motion = torch.cat((B_motion, scaled_logits), dim=1)

            scaled_logits_embedding = self.transformer.wte(scaled_logits)
            xs = torch.cat((x, scaled_logits_embedding), dim=1)
            

        return xs, B_motion

    def sample_for_eval_CFG_babel_inference(self, A_text, B_text, A_motion, if_categorial=False, length=6400, clip_model=None, device=torch.device('cuda'), tokenizer='clip', unit_length=4, reference_end_token=None, cfg=7.0, threshold=3):

        import clip
        B_token_length = length // unit_length - A_motion.shape[0]

        if tokenizer == 'clip':
            A_text = clip.tokenize(A_text, truncate=True).to(device)  
            A_feat_clip_text = clip_model.encode_text(A_text).float()  
            B_text = clip.tokenize(B_text, truncate=True).to(device)
            B_feat_clip_text = clip_model.encode_text(B_text).float()
        elif tokenizer == 't5-xxl':
            A_feat_clip_text = torch.from_numpy(clip_model.encode(A_text)).float()   
            A_feat_clip_text = A_feat_clip_text.to(device)
            B_feat_clip_text = torch.from_numpy(clip_model.encode(B_text)).float()
            B_feat_clip_text = B_feat_clip_text.to(device)
        
        A_text_embeddings = self.transformer.cond_embed(A_feat_clip_text).unsqueeze(0)
        A_text_embeddings = A_text_embeddings.unsqueeze(0)
        B_text_embeddings = self.transformer.cond_embed(B_feat_clip_text).unsqueeze(0)
        B_text_embeddings = B_text_embeddings.unsqueeze(0)

        A_motion = A_motion.unsqueeze(0)
        A_motion_embeddings = self.transformer.wte(A_motion)
        B_motion = torch.tensor([]).to(device)

        attention_weights = []

        for k in range(B_token_length): 
            if k == 0:    
                x = torch.cat([A_text_embeddings, A_motion_embeddings, B_text_embeddings], dim=1)
                
            else:
                x = xs

            
            
            conditions = self.forward_babel_eval(x, return_attention=False)
            conditions = conditions[:, -1, :]                            

            empty_clip_text = ''
            if tokenizer == 'clip':
                empty_text = clip.tokenize(empty_clip_text, truncate=True).to(device)
                empty_feat_clip_text = clip_model.encode_text(empty_text).float()
            elif tokenizer == 't5-xxl':
                empty_feat_clip_text = torch.from_numpy(clip_model.encode(empty_clip_text)).float()   
                empty_feat_clip_text = empty_feat_clip_text.unsqueeze(0)
                empty_feat_clip_text = empty_feat_clip_text.to(device)

            empty_feat_clip_text_embedding = self.transformer.cond_embed(empty_feat_clip_text).unsqueeze(0)
            
            if k == 0:    
                empty_input = torch.cat([empty_feat_clip_text_embedding, A_motion_embeddings, empty_feat_clip_text_embedding], dim=1)
                empty_conditions = self.forward_babel_eval(empty_input)
            else:
                B_motion_embeddings = self.transformer.wte(B_motion)
                empty_input = torch.cat([empty_feat_clip_text_embedding, A_motion_embeddings, empty_feat_clip_text_embedding, B_motion_embeddings], dim=1)
                empty_conditions = self.forward_babel_eval(empty_input)

            empty_conditions = empty_conditions[:, -1, :]                             
            temperature = 1.0
            
            mix_conditions = torch.cat([conditions, empty_conditions], dim=0)    
            sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg) 

            # chunk
            if cfg != 1:
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)        
            else:
                scaled_logits = sampled_token_latent

            scaled_logits = scaled_logits.unsqueeze(0)     
            
            if reference_end_token is not None:
                distance_l2 = torch.sqrt(torch.sum((scaled_logits - reference_end_token)**2))
                print(distance_l2)
                if distance_l2 < threshold:
                    break
            
            B_motion = torch.cat((B_motion, scaled_logits), dim=1)

            scaled_logits_embedding = self.transformer.wte(scaled_logits)
            xs = torch.cat((x, scaled_logits_embedding), dim=1)
            
        
        
        return xs, B_motion
    

    def sample_for_eval_CFG_babel_inference_new(self, B_text, A_motion, if_categorial=False, length=78, clip_model=None, device=torch.device('cuda'), tokenizer='clip', unit_length=4, reference_end_token=None, cfg=4.5, threshold=3):

        import clip 
        B_token_length = length // unit_length

        if tokenizer == 'clip':
            A_text = clip.tokenize(A_text, truncate=True).to(device)  
            A_feat_clip_text = clip_model.encode_text(A_text).float()  
            B_text = clip.tokenize(B_text, truncate=True).to(device)
            B_feat_clip_text = clip_model.encode_text(B_text).float()
        elif tokenizer == 't5-xxl':
            B_feat_clip_text = torch.from_numpy(clip_model.encode(B_text)).float()
            B_feat_clip_text = B_feat_clip_text.to(device)

        empty_clip_text = ''
        if tokenizer == 'clip':
            empty_text = clip.tokenize(empty_clip_text, truncate=True).to(device)
            empty_feat_clip_text = clip_model.encode_text(empty_text).float()
        elif tokenizer == 't5-xxl':
            empty_feat_clip_text = torch.from_numpy(clip_model.encode(empty_clip_text)).float()   
            empty_feat_clip_text = empty_feat_clip_text.unsqueeze(0)
            empty_feat_clip_text = empty_feat_clip_text.to(device)
        
        B_text_embeddings = self.transformer.cond_embed(B_feat_clip_text).unsqueeze(0)

        A_motion = A_motion.unsqueeze(0)
        A_motion_embeddings = self.transformer.wte(A_motion)
        B_motion = torch.tensor([]).to(device)

    
        attention_weights = []

        for k in range(B_token_length): 
            if k == 0: 
                x = torch.cat([B_text_embeddings, A_motion_embeddings], dim=1)
            else:
                x = xs

            conditions = self.forward_babel_eval(x, return_attention=False)
            conditions = conditions[:, -1, :]                             
            

            empty_feat_clip_text_embedding = self.transformer.cond_embed(empty_feat_clip_text).unsqueeze(0)
            
            if k == 0:
                empty_input = torch.cat([empty_feat_clip_text_embedding, A_motion_embeddings], dim=1)
                
                empty_conditions = self.forward_babel_eval(empty_input)
            else:
                B_motion_embeddings = self.transformer.wte(B_motion)
                empty_input = torch.cat([empty_feat_clip_text_embedding, A_motion_embeddings, B_motion_embeddings], dim=1)
                empty_conditions = self.forward_babel_eval(empty_input)

            empty_conditions = empty_conditions[:, -1, :]                             
            temperature = 1.0
           
            mix_conditions = torch.cat([conditions, empty_conditions], dim=0)     
            sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg) 

            # chunk
            if cfg != 1:
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)       
            else:
                scaled_logits = sampled_token_latent

            scaled_logits = scaled_logits.unsqueeze(0)      
            
            if reference_end_token is not None:
                distance_l2 = torch.sqrt(torch.sum((scaled_logits - reference_end_token)**2))
                print(distance_l2)
                if distance_l2 < threshold:
                    break
            
            B_motion = torch.cat((B_motion, scaled_logits), dim=1)

            scaled_logits_embedding = self.transformer.wte(scaled_logits)
            xs = torch.cat((x, scaled_logits_embedding), dim=1)
            
        
        
        return xs, B_motion


    def sample_for_eval_CFG_babel_inference_new_demo(self, B_text, A_motion, if_categorial=False, length=312, clip_model=None, device=torch.device('cuda'), tokenizer='clip', unit_length=4, reference_end_token=None, cfg=4.5, threshold=3, temperature=1.0):

        import clip 
        B_token_length = length // unit_length - A_motion.shape[0]
        
        if tokenizer == 'clip':
            A_text = clip.tokenize(A_text, truncate=True).to(device) 
            A_feat_clip_text = clip_model.encode_text(A_text).float()  
            B_text = clip.tokenize(B_text, truncate=True).to(device)
            B_feat_clip_text = clip_model.encode_text(B_text).float()
        elif tokenizer == 't5-xxl':
            B_feat_clip_text = torch.from_numpy(clip_model.encode(B_text)).float()
            B_feat_clip_text = B_feat_clip_text.to(device)

        empty_clip_text = ''
        if tokenizer == 'clip':
            empty_text = clip.tokenize(empty_clip_text, truncate=True).to(device)
            empty_feat_clip_text = clip_model.encode_text(empty_text).float()
        elif tokenizer == 't5-xxl':
            empty_feat_clip_text = torch.from_numpy(clip_model.encode(empty_clip_text)).float()   
            empty_feat_clip_text = empty_feat_clip_text.unsqueeze(0)
            empty_feat_clip_text = empty_feat_clip_text.to(device)
        
        B_text_embeddings = self.transformer.cond_embed(B_feat_clip_text).unsqueeze(0)
        B_text_embeddings = B_text_embeddings.unsqueeze(0)

        A_motion = A_motion.unsqueeze(0)
        A_motion_embeddings = self.transformer.wte(A_motion)
        B_motion = torch.tensor([]).to(device)

        # 存储所有层的注意力权重
        attention_weights = []

        for k in range(B_token_length): 
            if k == 0: 
                x = torch.cat([B_text_embeddings, A_motion_embeddings], dim=1)
               
            else:
                x = xs

            
            conditions = self.forward_babel_eval(x, return_attention=False)
            conditions = conditions[:, -1, :]                             
            

            empty_feat_clip_text_embedding = self.transformer.cond_embed(empty_feat_clip_text).unsqueeze(0)
            
            if k == 0:
                empty_input = torch.cat([empty_feat_clip_text_embedding, A_motion_embeddings], dim=1)
                empty_conditions = self.forward_babel_eval(empty_input)
            else:
                B_motion_embeddings = self.transformer.wte(B_motion)
                empty_input = torch.cat([empty_feat_clip_text_embedding, A_motion_embeddings, B_motion_embeddings], dim=1)
                empty_conditions = self.forward_babel_eval(empty_input)

            empty_conditions = empty_conditions[:, -1, :]                             
           
            mix_conditions = torch.cat([conditions, empty_conditions], dim=0)     
            sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg)  

            # chunk
            if cfg != 1:
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)           
            else:
                scaled_logits = sampled_token_latent

            scaled_logits = scaled_logits.unsqueeze(0)      
            
            if reference_end_token is not None:
                distance_l2 = torch.sqrt(torch.sum((scaled_logits - reference_end_token)**2))
                print(distance_l2)
                if distance_l2 < threshold and k > 10:
                    break
            
            B_motion = torch.cat((B_motion, scaled_logits), dim=1)

            scaled_logits_embedding = self.transformer.wte(scaled_logits)
            xs = torch.cat((x, scaled_logits_embedding), dim=1)
            
        
        
        return xs, B_motion

    def sample_for_eval_CFG_babel_inference_two_forward(self, B_text, A_motion, if_categorial=False, length=312, clip_model=None, device=torch.device('cuda'), tokenizer='clip', unit_length=4, reference_end_token=None, cfg=4.5, threshold=3, temperature=1.0):
        """
        Inference loop that mimics the "Two-Forward" training strategy.
        This version is rewritten for clarity and to fix dimension mismatches.
        """
        import clip
        print("\n--- [DEBUG] Entering Two-Forward Inference ---")

        B_token_length = length // unit_length - A_motion.shape[0]

        if tokenizer == 't5-xxl':
            B_feat_clip_text = torch.from_numpy(clip_model.encode(B_text)).float().to(device)
        else:
            raise NotImplementedError("Only t5-xxl is supported for this function.")
        empty_feat_clip_text = torch.from_numpy(clip_model.encode('')).float().unsqueeze(0).to(device)

        # --- Consistently create 3D embeddings [batch, seq, dim] ---
        B_text_embeddings = self.transformer.cond_embed(B_feat_clip_text).unsqueeze(0).unsqueeze(0)
        empty_text_embeddings = self.transformer.cond_embed(empty_feat_clip_text).unsqueeze(0)
        print(f"[DEBUG] Initial B_text_embeddings shape: {B_text_embeddings.shape}")
        print(f"[DEBUG] Initial empty_text_embeddings shape: {empty_text_embeddings.shape}")

        # --- Initial motion history setup ---
        A_motion_embeddings = self.transformer.wte(A_motion.unsqueeze(0))
        current_motion_tokens = A_motion
        print(f"[DEBUG] Initial A_motion_embeddings shape: {A_motion_embeddings.shape}")
        print(f"[DEBUG] Initial current_motion_tokens shape: {current_motion_tokens.shape}")

        for k in range(B_token_length):
            print(f"\n--- [DEBUG] Generating Token {k+1}/{B_token_length} ---")

            # --- Prepare motion embeddings for this step ---
            if k == 0:
                # Use the empty initial motion history
                step_motion_embeddings = A_motion_embeddings
            else:
                # Embed the current token history
                step_motion_embeddings = self.transformer.wte(current_motion_tokens.unsqueeze(0))
            print(f"[DEBUG] step_motion_embeddings shape: {step_motion_embeddings.shape}")

            # === 1. First Forward Pass (Rough Draft) ===
            x_first_pass = torch.cat([B_text_embeddings, step_motion_embeddings], dim=1)
            empty_input_first_pass = torch.cat([empty_text_embeddings, step_motion_embeddings], dim=1)
            print(f"[DEBUG] x_first_pass shape: {x_first_pass.shape}")
            print(f"[DEBUG] empty_input_first_pass shape: {empty_input_first_pass.shape}")

            # Get conditions from the transformer
            conditions_first_pass = self.forward_babel_eval(x_first_pass, return_attention=False)[:, -1, :]
            empty_conditions_first_pass = self.forward_babel_eval(empty_input_first_pass)[:, -1, :]

            # Sample a rough prediction for the next token using diffusion
            mix_conditions_first_pass = torch.cat([conditions_first_pass, empty_conditions_first_pass], dim=0)
            pred_xstart_rough = self.diff_loss.sample(mix_conditions_first_pass, temperature=temperature, cfg=cfg)
            if cfg != 1:
                pred_xstart_rough, _ = pred_xstart_rough.chunk(2, dim=0)
            print(f"[DEBUG] pred_xstart_rough shape: {pred_xstart_rough.shape}")

            # === 2. Second Forward Pass (Refined) ===
            # Create the "dirtied" input by appending the rough prediction
            updated_motion_tokens = torch.cat([current_motion_tokens, pred_xstart_rough], dim=0)
            updated_motion_embeddings = self.transformer.wte(updated_motion_tokens.unsqueeze(0))
            print(f"[DEBUG] updated_motion_embeddings shape: {updated_motion_embeddings.shape}")

            x_second_pass = torch.cat([B_text_embeddings, updated_motion_embeddings], dim=1)
            empty_input_second_pass = torch.cat([empty_text_embeddings, updated_motion_embeddings], dim=1)

            conditions_second_pass = self.forward_babel_eval(x_second_pass, return_attention=False)[:, -1, :]
            empty_conditions_second_pass = self.forward_babel_eval(empty_input_second_pass)[:, -1, :]

            # Sample the final, refined token using diffusion
            mix_conditions_second_pass = torch.cat([conditions_second_pass, empty_conditions_second_pass], dim=0)
            final_token, _ = self.diff_loss.sample(mix_conditions_second_pass, temperature=temperature, cfg=cfg).chunk(2, dim=0)
            print(f"[DEBUG] final_token shape: {final_token.shape}")

            # === 3. Update History ===
            current_motion_tokens = torch.cat([current_motion_tokens, final_token], dim=0)
            print(f"[DEBUG] New current_motion_tokens shape: {current_motion_tokens.shape}")

        print("\n--- [DEBUG] Finished Token Generation ---")
        # Return only the newly generated tokens (B_motion)
        B_motion = current_motion_tokens[A_motion.shape[0]:, :].unsqueeze(0)
        return None, B_motion

    
    #--------------Test classification head--------------------
    def sample_for_eval_classification(self, clip_text, if_categorial=False, length=196, clip_model=None, device=torch.device('cuda'), tokenizer='clip', unit_length=4):
        
        import clip
       
        
        for k in range(51): 
            if k == 0:
                x = []
            else:
                x = xs

            if tokenizer == 'clip':
                text = clip.tokenize(clip_text, truncate=True).to(device)  

                feat_clip_text = clip_model.encode_text(text).float() 
            elif tokenizer == 't5-xxl':
                feat_clip_text = torch.from_numpy(clip_model.module.encode(clip_text)).float()  
           
            conditions = self.forward(x, feat_clip_text)
            conditions = conditions[:, -1, :]     

            empty_clip_text = ''
            if tokenizer == 'clip':
                empty_text = clip.tokenize(empty_clip_text, truncate=True).to(device)
                empty_feat_clip_text = clip_model.encode_text(empty_text).float()
            elif tokenizer == 't5-xxl':
                empty_feat_clip_text = torch.from_numpy(clip_model.module.encode(empty_clip_text)).float()  
                empty_feat_clip_text = empty_feat_clip_text.unsqueeze(0)
                empty_feat_clip_text = empty_feat_clip_text.to(device)

            empty_conditions = self.forward(x, empty_feat_clip_text)     
            empty_conditions = empty_conditions[:, -1, :]                    
           
            temperature = 1.0
            cfg = 7.5
            
            mix_conditions = torch.cat([conditions, empty_conditions], dim=0) 
            sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg)  

            # chunk
            if cfg != 1:
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)        
            else:
                scaled_logits = sampled_token_latent


            prediction_logits = self.classify_head(conditions)  
            probs = torch.sigmoid(prediction_logits)
            predicted_classes = torch.argmax(probs, dim=-1)
            

            scaled_logits = scaled_logits.unsqueeze(0)      
                
            if k == 0:
                xs = scaled_logits
            else:
                xs = torch.cat((xs, scaled_logits), dim=1)

            if predicted_classes == 1:
                break

        return xs
    

    #--------------------Test CFG-----------------------
    def sample_for_eval_CFG_test(self, clip_text, if_categorial=False, length=196, clip_model=None, cfg=1, device=torch.device('cuda'), tokenizer='clip', unit_length=4):

        import clip
        max_token_len = length // unit_length
        
        
        for k in range(max_token_len): 
            if k == 0:
                x = []
            else:
                x = xs

            
            if cfg != 1:
                if tokenizer == 'clip':
                    text = clip.tokenize(clip_text, truncate=True).to(device)  

                    feat_clip_text = clip_model.encode_text(text).float() 
                elif tokenizer == 't5-xxl':
                    feat_clip_text = torch.from_numpy(clip_model.module.encode(clip_text)).float()  
                
                conditions = self.forward(x, feat_clip_text)       
                
                conditions = conditions[:, -1, :]                            
                empty_clip_text = ''
                if tokenizer == 'clip':
                    empty_text = clip.tokenize(empty_clip_text, truncate=True).to(device)
                    empty_feat_clip_text = clip_model.encode_text(empty_text).float()
                elif tokenizer == 't5-xxl':
                    empty_feat_clip_text = torch.from_numpy(clip_model.module.encode(empty_clip_text)).float()   
                    empty_feat_clip_text = empty_feat_clip_text.unsqueeze(0)
                    empty_feat_clip_text = empty_feat_clip_text.to(device)

                empty_conditions = self.forward(x, empty_feat_clip_text)     
                empty_conditions = empty_conditions[:, -1, :]                             
                temperature = 1.0
                
                
                mix_conditions = torch.cat([conditions, empty_conditions], dim=0)     
                sampled_token_latent = self.diff_loss.sample(mix_conditions, temperature=temperature, cfg=cfg)  

                # chunk
                scaled_logits, _ = sampled_token_latent.chunk(2, dim=0)           
            
            else:   
                if tokenizer == 'clip':
                    text = clip.tokenize(clip_text, truncate=True).to(device)  
                    feat_clip_text = clip_model.encode_text(text).float() 
                elif tokenizer == 't5-xxl':
                    feat_clip_text = torch.from_numpy(clip_model.module.encode(clip_text)).float()  
                    feat_clip_text = feat_clip_text.to(device)


                conditions = self.forward(x, feat_clip_text)       
  
                conditions = conditions[:, -1, :]                             
                temperature = 1.0
                sampled_token_latent = self.diff_loss.sample(conditions, temperature=temperature, cfg=cfg) 
                scaled_logits = sampled_token_latent       

            scaled_logits = scaled_logits.unsqueeze(0)      
                
            if k == 0:
                xs = scaled_logits
            else:
                xs = torch.cat((xs, scaled_logits), dim=1)

        return xs
    #--------------------------------------------------

    def forward_discrete(self, idx: torch.Tensor, clip_feature: torch.Tensor, use_cache=False, past_key_values=None) -> torch.Tensor:
        if len(idx) == 0:
            token_embeddings = self.transformer.cond_embed(clip_feature).unsqueeze(0)

        else:
            b, t = idx.size()
            #idx = idx.float()
            assert (
                t <= self.config.block_size
            ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"

            # forward the LLaMA model itself
            token_embeddings = self.transformer.wte(idx)  
            text_embeddings = self.transformer.cond_embed(clip_feature).unsqueeze(1)        
            token_embeddings = torch.cat([text_embeddings, token_embeddings], dim=1) 
        
        x = token_embeddings

        # -------------------kv cache-------------------
        #presents = () if use_cache else None
        if use_cache:
            if past_key_values is None:
                past_key_values = [None] * len(self.transformer.h)   
            

        for i,block in enumerate(self.transformer.h):
            if use_cache:
                last_past = past_key_values[i]     
                x, presents = block(x, last_past, use_cache)
                past_key_values[i] = list(presents)
            else:
                x = block(x)
        x = self.transformer.ln_f(x)

        logits = self.lm_head(x) 
        

        return logits


    def forward(self, idx: torch.Tensor, feature: torch.Tensor) -> torch.Tensor:
        if len(idx) == 0:
            token_embeddings = self.transformer.cond_embed(feature).unsqueeze(0)

        else:
            b, t, c = idx.size()
            idx = idx.float()
            assert (
                t <= self.config.block_size
            ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"

            # forward the LLaMA model itself
            token_embeddings = self.transformer.wte(idx)  
            text_embeddings = self.transformer.cond_embed(feature).unsqueeze(1)    
            token_embeddings = torch.cat([text_embeddings, token_embeddings], dim=1) 
        
        x = token_embeddings  

        for i,block in enumerate(self.transformer.h):
            x = block(x)
        x = self.transformer.ln_f(x)    
        logits = self.out_proj(x)
        return logits


    def forward_inference(self, idx: torch.Tensor, feature: torch.Tensor) -> torch.Tensor:
        if len(idx) == 0:
            token_embeddings = self.transformer.cond_embed(feature).unsqueeze(0)

        else:
            b, t, c = idx.size()
            idx = idx.float()
            assert (
                t <= self.config.block_size
            ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"

            # forward the LLaMA model itself
            token_embeddings = self.transformer.wte(idx)  
            text_embeddings = self.transformer.cond_embed(feature).unsqueeze(0)    
            token_embeddings = torch.cat([text_embeddings.unsqueeze(0), token_embeddings], dim=1) 
        
        x = token_embeddings  

        if len(x.shape) == 2:
            x = x.unsqueeze(0)

        for i,block in enumerate(self.transformer.h):
            x = block(x)
        x = self.transformer.ln_f(x)    
        logits = self.out_proj(x)
        return logits
        

    def babel_long(self, idx: torch.Tensor, clip_feature: torch.Tensor, use_cache=False, past_key_values=None, num_subseq=None, length=None) -> torch.Tensor:
        
        b, t, c = idx.size()
        idx = idx.float()
        idx = self.transformer.wte(idx)
        assert (
                t <= self.config.block_size
            ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
        for i in range(b):
            length_i = length[i][:num_subseq[i]]       
            clip_feature_i = clip_feature[i][:num_subseq[i]]

            pointer = 0
            for j in range(num_subseq[i]):
                if j > 0:
                    pointer += length_i[j].item()
                    pointer += 1
                pointer = int(pointer)

                clip_feature_i_j = self.transformer.cond_embed(clip_feature_i[j].unsqueeze(0)).unsqueeze(1)
                idx[i] = torch.cat([idx[i][:pointer].unsqueeze(0), clip_feature_i_j, idx[i][pointer:-1].unsqueeze(0)], dim=1)[0]
                
        x = idx


        if use_cache:
            if past_key_values is None:
                past_key_values = [None] * len(self.transformer.h)   
            

        for i,block in enumerate(self.transformer.h):
            if use_cache:
                last_past = past_key_values[i]      
                x, presents = block(x, last_past, use_cache)
                past_key_values[i] = list(presents)
            else:
                x = block(x)
        x = self.transformer.ln_f(x)

        logits = self.out_proj(x)
        return logits
        

    def forward_babel_eval(self, x, return_attention=False) -> torch.Tensor:
        layer_attentions = []
        for block in self.transformer.h:
            if return_attention:
                x, att = block(x, return_attention=True)
                layer_attentions.append(att)
            else:
                x = block(x)
        
        x = self.transformer.ln_f(x)
        if self.use_out_proj:
            logits = self.out_proj(x)
        else:
            logits = x
        
        if return_attention:
            return logits, layer_attentions
        return logits
    
    def forward_babel(self, idx: torch.Tensor, clip_feature: torch.Tensor, A_token_length) -> torch.Tensor:
        if len(idx) == 0:   # inference
            token_embeddings = self.transformer.cond_embed(clip_feature).unsqueeze(1)

        else:
            b, t, c = idx.size()
            idx = idx.float()
            assert (
                t <= self.config.block_size
            ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"

            

            A_feature = clip_feature[:, 0, :]             
            B_feature = clip_feature[:, 1, :]                   


            A_text_embeddings = self.transformer.cond_embed(A_feature).unsqueeze(1)
            B_text_embeddings = self.transformer.cond_embed(B_feature).unsqueeze(1)

            token_embeddings = torch.zeros(b, self.config.block_size, self.config.n_embd).to(idx.device)
            for i in range(b):
                A_idx = idx[i, :A_token_length[i].item(), :]
                B_idx = idx[i, A_token_length[i].item():-2, :]
                token_embeddings[i, :, :] = torch.cat([A_text_embeddings[i], self.BOM_tag, self.transformer.wte(A_idx), B_text_embeddings[i], self.BOM_tag, self.transformer.wte(B_idx)], dim=0)  #token_embeddings.shape = (b,t+1,1024)
        
        x = token_embeddings
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)

        if self.use_out_proj:
            logits = self.out_proj(x)
        else:
            logits = x


        return logits

    def forward_babel2(self, idx: torch.Tensor, clip_feature: torch.Tensor) -> torch.Tensor:
        if len(idx) == 0:   # inference
            token_embeddings = self.transformer.cond_embed(clip_feature).unsqueeze(1)

        else:
            b, t, c = idx.size()
            idx = idx.float()
            assert (
                t <= self.config.block_size
            ), f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"

            B_feature = clip_feature
            B_text_embeddings = self.transformer.cond_embed(B_feature)

            idx_embeddings = self.transformer.wte(idx)

            
            token_embeddings = torch.cat([B_text_embeddings, idx_embeddings], dim=1)
            
        
        x = token_embeddings
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)

        if self.use_out_proj:
            logits = self.out_proj(x)
        else:
            logits = x

        return logits
    

    def resize_token_embeddings(
        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None, using_old_initilization: bool = False
    ) -> nn.Embedding:
        """
        Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`.

        Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.

        Arguments:
            new_num_tokens (`int`, *optional*):
                The new number of tokens in the embedding matrix. Increasing the size will add newly initialized
                vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just
                returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything.
            pad_to_multiple_of (`int`, *optional*):
                If set will pad the embedding matrix to a multiple of the provided value.If `new_num_tokens` is set to
                `None` will just pad the embedding to a multiple of `pad_to_multiple_of`.

                This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
                `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. For more
                details about this, or help on choosing the correct value for resizing, refer to this guide:
                https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc

        Return:
            `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model.
        """
        model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
        if new_num_tokens is None and pad_to_multiple_of is None:
            return model_embeds

        # Update base model and current model config
        self.config.vocab_size = model_embeds.weight.shape[0]
        self.vocab_size = model_embeds.weight.shape[0]

        # Tie weights again if needed
        # self.tie_weights()

        return model_embeds
    
    def _resize_token_embeddings(self, new_num_tokens, pad_to_multiple_of=None):
        old_embeddings = self.get_input_embeddings()
        new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens, pad_to_multiple_of)
        old_embeddings_requires_grad = old_embeddings.weight.requires_grad
        new_embeddings.requires_grad_(old_embeddings_requires_grad)
        self.set_input_embeddings(new_embeddings)

        # Update new_num_tokens with the actual size of new_embeddings
        if pad_to_multiple_of is not None:
            # if is_deepspeed_zero3_enabled():
            #     import deepspeed

            #     with deepspeed.zero.GatheredParameters(new_embeddings.weight, modifier_rank=None):
            #         new_num_tokens = new_embeddings.weight.shape[0]
            # else:
            new_num_tokens = new_embeddings.weight.shape[0]

        # if word embeddings are not tied, make sure that lm head is resized as well
        # if self.get_output_embeddings() is not None and not self.config.tie_word_embeddings:
        if self.get_output_embeddings() is not None and not False:
            old_lm_head = self.get_output_embeddings()
            new_lm_head = self._get_resized_lm_head(old_lm_head, new_num_tokens)
            # if hasattr(old_lm_head, "_hf_hook"):
            #     hook = old_lm_head._hf_hook
            #     add_hook_to_module(new_lm_head, hook)
            old_lm_head_requires_grad = old_lm_head.weight.requires_grad
            new_lm_head.requires_grad_(old_lm_head_requires_grad)
            self.set_output_embeddings(new_lm_head)

        return self.get_input_embeddings()
    
    def _get_resized_embeddings(
        self,
        old_embeddings: nn.Embedding,
        new_num_tokens: Optional[int] = None,
        pad_to_multiple_of: Optional[int] = None,
    ) -> nn.Embedding:
        """
        Build a resized Embedding Module from a provided token Embedding Module. Increasing the size will add newly
        initialized vectors at the end. Reducing the size will remove vectors from the end

        Args:
            old_embeddings (`torch.nn.Embedding`):
                Old embeddings to be resized.
            new_num_tokens (`int`, *optional*):
                New number of tokens in the embedding matrix.

                Increasing the size will add newly initialized vectors at the end. Reducing the size will remove
                vectors from the end. If not provided or `None`, just returns a pointer to the input tokens
                `torch.nn.Embedding` module of the model without doing anything.
            pad_to_multiple_of (`int`, *optional*):
                If set will pad the embedding matrix to a multiple of the provided value. If `new_num_tokens` is set to
                `None` will just pad the embedding to a multiple of `pad_to_multiple_of`.

                This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
                `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. For more
                details about this, or help on choosing the correct value for resizing, refer to this guide:
                https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc


        Return:
            `torch.nn.Embedding`: Pointer to the resized Embedding Module or the old Embedding Module if
            `new_num_tokens` is `None`
        """

        if pad_to_multiple_of is not None:
            if not isinstance(pad_to_multiple_of, int):
                raise ValueError(
                    f"Asking to pad the embedding matrix to a multiple of `{pad_to_multiple_of}`, which is not and integer. Please make sure to pass an integer"
                )
            if new_num_tokens is None:
                new_num_tokens = old_embeddings.weight.shape[0]
            new_num_tokens = ((new_num_tokens + pad_to_multiple_of - 1) // pad_to_multiple_of) * pad_to_multiple_of
        else:
            print(
                "You are resizing the embedding layer without providing a `pad_to_multiple_of` parameter. This means that the new embedding"
                f" dimension will be {new_num_tokens}. This might induce some performance reduction as *Tensor Cores* will not be available."
                " For more details about this, or help on choosing the correct value for resizing, refer to this guide:"
                " https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc"
            )

        if new_num_tokens is None:
            return old_embeddings

        # if is_deepspeed_zero3_enabled():
        if False:
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_embeddings.weight, modifier_rank=None):
                old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
        else:
            old_num_tokens, old_embedding_dim = old_embeddings.weight.size()

        # if old_num_tokens == new_num_tokens and not is_deepspeed_zero3_enabled():
        if old_num_tokens == new_num_tokens and not False:
            return old_embeddings

        if not isinstance(old_embeddings, nn.Embedding):
            raise TypeError(
                f"Old embeddings are of type {type(old_embeddings)}, which is not an instance of {nn.Embedding}. You"
                " should either use a different resize function or make sure that `old_embeddings` are an instance of"
                f" {nn.Embedding}."
            )

        # Build new embeddings

        # When using DeepSpeed ZeRO-3, we shouldn't create new embeddings with DeepSpeed init
        # because the shape of the new embedding layer is used across various modeling files
        # as well as to update config vocab size. Shape will be 0 when using DeepSpeed init leading
        # to errors when training.
        new_embeddings = nn.Embedding(
            new_num_tokens,
            old_embedding_dim,
            device=old_embeddings.weight.device,
            dtype=old_embeddings.weight.dtype,
        )

        # initialize all new embeddings (in particular added tokens)
        self._init_weights(new_embeddings)

        # Copy token embeddings from the previous weights

        # numbers of tokens to copy
        n = min(old_num_tokens, new_num_tokens)

        # if is_deepspeed_zero3_enabled():
        if False:
            import deepspeed

            params = [old_embeddings.weight, new_embeddings.weight]
            with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
                new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :]
        else:
            new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :]

        return new_embeddings


    def _get_resized_lm_head(
        self, old_lm_head: nn.Linear, new_num_tokens: Optional[int] = None, transposed: Optional[bool] = False
    ) -> nn.Linear:
        """
        Build a resized Linear Module from a provided old Linear Module. Increasing the size will add newly initialized
        vectors at the end. Reducing the size will remove vectors from the end

        Args:
            old_lm_head (`torch.nn.Linear`):
                Old lm head liner layer to be resized.
            new_num_tokens (`int`, *optional*):
                New number of tokens in the linear matrix.

                Increasing the size will add newly initialized vectors at the end. Reducing the size will remove
                vectors from the end. If not provided or `None`, just returns a pointer to the input tokens
                `torch.nn.Linear` module of the model without doing anything. transposed (`bool`, *optional*, defaults
                to `False`): Whether `old_lm_head` is transposed or not. If True `old_lm_head.size()` is `lm_head_dim,
                vocab_size` else `vocab_size, lm_head_dim`.

        Return:
            `torch.nn.Linear`: Pointer to the resized Linear Module or the old Linear Module if `new_num_tokens` is
            `None`
        """
        if new_num_tokens is None:
            return old_lm_head

        # if is_deepspeed_zero3_enabled():
        if False:
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_lm_head.weight, modifier_rank=None):
                old_num_tokens, old_lm_head_dim = (
                    old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size()
                )
        else:
            old_num_tokens, old_lm_head_dim = (
                old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size()
            )

        # if old_num_tokens == new_num_tokens and not is_deepspeed_zero3_enabled():
        if old_num_tokens == new_num_tokens and not False:
            return old_lm_head

        if not isinstance(old_lm_head, nn.Linear):
            raise TypeError(
                f"Old language model head is of type {type(old_lm_head)}, which is not an instance of {nn.Linear}. You"
                " should either use a different resize function or make sure that `old_lm_head` are an instance of"
                f" {nn.Linear}."
            )

        # Build new lm head
        new_lm_head_shape = (old_lm_head_dim, new_num_tokens) if not transposed else (new_num_tokens, old_lm_head_dim)
        has_new_lm_head_bias = old_lm_head.bias is not None

        # When using DeepSpeed ZeRO-3, we shouldn't create new embeddings with DeepSpeed init
        # because the shape of the new embedding layer is used across various modeling files
        # as well as to update config vocab size. Shape will be 0 when using DeepSpeed init leading
        # to errors when training.
        new_lm_head = nn.Linear(
            *new_lm_head_shape,
            bias=has_new_lm_head_bias,
            device=old_lm_head.weight.device,
            dtype=old_lm_head.weight.dtype,
        )

        # initialize new lm head (in particular added tokens)
        self._init_weights(new_lm_head)

        num_tokens_to_copy = min(old_num_tokens, new_num_tokens)

        # if is_deepspeed_zero3_enabled():
        if False:
            import deepspeed

            params = [old_lm_head.weight, old_lm_head.bias, new_lm_head.weight, new_lm_head.bias]
            with deepspeed.zero.GatheredParameters(params, modifier_rank=0):
                self._copy_lm_head_original_to_resized(
                    new_lm_head, old_lm_head, num_tokens_to_copy, transposed, has_new_lm_head_bias
                )
        else:
            self._copy_lm_head_original_to_resized(
                new_lm_head, old_lm_head, num_tokens_to_copy, transposed, has_new_lm_head_bias
            )

        return new_lm_head

    def _copy_lm_head_original_to_resized(
        self, new_lm_head, old_lm_head, num_tokens_to_copy, transposed, has_new_lm_head_bias
    ):
        # Copy old lm head weights to new lm head
        if not transposed:
            new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[:num_tokens_to_copy, :]
        else:
            new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[:, :num_tokens_to_copy]

        # Copy bias weights to new lm head
        if has_new_lm_head_bias:
            new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy]

    @classmethod
    def from_name(cls, name: str) -> Self:
        return cls(LLaMAHFConfig.from_name(name))


class Block(nn.Module):
    def __init__(self, config: LLaMAHFConfig) -> None:
        super().__init__()
        self.rms_1 = RMSNorm(config.n_embd)

        # sentence level:
        self.attn = CausalSelfAttention(config)        
        self.rms_2 = RMSNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x: torch.Tensor, last_past=None, use_cache=False, return_attention=False) -> torch.Tensor:
        if use_cache:  
            if return_attention:
                a, attn = self.attn.forward_attn(self.rms_1(x), last_past, use_cache)
            else:
                a, present = self.attn(self.rms_1(x), last_past, use_cache)
            x = x + a
        else:
            if return_attention:
                a, attn = self.attn.forward_attn(self.rms_1(x))
            else:
                a = self.attn(self.rms_1(x))
            x = x + a
        x = x + self.mlp(self.rms_2(x))

        if use_cache:
            if return_attention:
                return x, present, attn
            else:
                return x, present  
        else:
            if return_attention:
                return x, attn
            else:
                return x


class CausalSelfAttention(nn.Module):
    def __init__(self, config: LLaMAHFConfig) -> None:
        super().__init__()
        assert config.n_embd % config.n_head == 0

        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)

        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.block_size = config.block_size
        self.rope_cache = None

        def scaling_factor(sequence_threshold):
            return np.log2((sequence_threshold**2) - sequence_threshold)
        scale_init = scaling_factor(self.block_size) 
        self.scale = nn.Parameter(torch.tensor(scale_init))

    def forward(self, x: torch.Tensor, last_past=None, use_cache=False) -> torch.Tensor:
        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)

        head_size = C // self.n_head
        k = k.view(B, T, self.n_head, head_size).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, head_size).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, head_size).transpose(1, 2)  # (B, nh, T, hs)

        # kv_cache
        if use_cache:
            if last_past is not None:
                past_key, past_value = last_past
                k = torch.cat([past_key, k], dim=-2)  
                v = torch.cat([past_value, v], dim=-2)   
            # else:
            #     key_states = k
            #     value_states = v

        if use_cache:
            present = (k, v)
        else:
            present = None

        # QK-Norm
        q = F.normalize(q, p=2, dim=-1)
        k = F.normalize(k, p=2, dim=-1)

        if self.rope_cache is None:
            # cache for future forward calls
            self.rope_cache = build_rope_cache(
                seq_len=self.block_size,
                n_elem=self.n_embd // self.n_head, 
                dtype=x.dtype,
                device=x.device,
            )


        q = apply_rope(q, self.rope_cache)
        k = apply_rope(k, self.rope_cache)


        
        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        #  att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        #  att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
        #  att = F.softmax(att, dim=-1)
        #  y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)

        # efficient attention using Flash Attention CUDA kernels
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, scale=self.scale.item())
        
        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side

        # output projection
        y = self.c_proj(y)

        
        if use_cache:
            return y, present
        return y

    def forward_attn(self, x: torch.Tensor, last_past=None, use_cache=False) -> torch.Tensor:
        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)

        head_size = C // self.n_head
        k = k.view(B, T, self.n_head, head_size).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, head_size).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, head_size).transpose(1, 2)  # (B, nh, T, hs)

        # kv_cache
        if use_cache:
            if last_past is not None:
                past_key, past_value = last_past
                k = torch.cat([past_key, k], dim=-2)  
                v = torch.cat([past_value, v], dim=-2) 
            # else:
            #     key_states = k
            #     value_states = v

        if use_cache:
            present = (k, v)
        else:
            present = None

        # QK-Norm
        q = F.normalize(q, p=2, dim=-1)
        k = F.normalize(k, p=2, dim=-1)

        if self.rope_cache is None:
            # cache for future forward calls
            self.rope_cache = build_rope_cache(
                seq_len=self.block_size,
                n_elem=self.n_embd // self.n_head, 
                dtype=x.dtype,
                device=x.device,
            )


        q = apply_rope(q, self.rope_cache)
        k = apply_rope(k, self.rope_cache)


        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = F.softmax(att, dim=-1)  # [B, n_head, T, T]

        # efficient attention using Flash Attention CUDA kernels
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side

        # output projection
        y = self.c_proj(y)
        
        return y, att

class LengthCausalSelfAttention(nn.Module):
    def __init__(self, config: LLaMAHFConfig) -> None:
        super().__init__()
        assert config.n_embd % config.n_head == 0

        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)

        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.block_size = config.block_size
        self.rope_cache = None

    def forward(self, x: torch.Tensor, y_mask: torch.Tensor) -> torch.Tensor:
        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)

        head_size = C // self.n_head
        k = k.view(B, T, self.n_head, head_size).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, head_size).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, head_size).transpose(1, 2)  # (B, nh, T, hs)

        if self.rope_cache is None:
            # cache for future forward calls
            self.rope_cache = build_rope_cache(
                seq_len=self.block_size,
                n_elem=self.n_embd // self.n_head, 
                dtype=x.dtype,
                device=x.device,
            )


        # q: 1, 16, 40 ,64
        # q: 128, 16, 106, 64
        q = apply_rope(q, self.rope_cache)
        k = apply_rope(k, self.rope_cache)
 
        attn_mask = torch.ones(T, T, dtype=torch.bool, device=x.device)
        attn_mask = torch.tril(attn_mask) 
        attn_mask = attn_mask.unsqueeze(0).expand(B, -1, -1)

        text_mask = y_mask.unsqueeze(2)*y_mask.unsqueeze(1)
        text_mask = F.pad(text_mask, (0, T-y_mask.shape[1], 0, T-y_mask.shape[1]), mode='constant', value=0)
        attn_mask = torch.logical_or(attn_mask, text_mask)

        y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask.unsqueeze(1), dropout_p=0.0, is_causal=False)

        y = y.transpose(1, 2).contiguous().view(B, T, C) 

     
        y = self.c_proj(y)

        return y


class MLP(nn.Module):
    def __init__(self, config: LLaMAHFConfig) -> None:
        super().__init__()
        hidden_dim = 4 * config.n_embd
        n_hidden = int(2 * hidden_dim / 3)
        N = 256
        # ensure n_hidden is multiple of N
        n_hidden = ((n_hidden - 1) // N) * N + N

        self.c_fc1 = nn.Linear(config.n_embd, n_hidden, bias=False)
        self.c_fc2 = nn.Linear(config.n_embd, n_hidden, bias=False)
        self.c_proj = nn.Linear(n_hidden, config.n_embd, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
       
        x = F.silu(self.c_fc1(x)) * self.c_fc2(x)
        x = self.c_proj(x)
        return x


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization.

    Derived from https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py. BSD 3-Clause License:
    https://github.com/bzhangGo/rmsnorm/blob/master/LICENSE.
    """

    def __init__(self, size: int, dim: int = -1, eps: float = 1e-5) -> None:
        super().__init__()
        self.scale = nn.Parameter(torch.ones(size))
        self.eps = eps
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # NOTE: the original RMSNorm paper implementation is not equivalent
        # norm_x = x.norm(2, dim=self.dim, keepdim=True)
        # rms_x = norm_x * d_x ** (-1. / 2)
        # x_normed = x / (rms_x + self.eps)
        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
        x_normed = x * torch.rsqrt(norm_x + self.eps)
        return self.scale * x_normed


def build_rope_cache(seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000) -> torch.Tensor:
    """Enhanced Transformer with Rotary Position Embedding.

    Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
    transformers/rope/__init__.py. MIT License:
    https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
    """
    # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))

    # Create position indexes `[0, 1, ..., seq_len - 1]`
    seq_idx = torch.arange(seq_len, dtype=dtype, device=device)

    # Calculate the product of position index and $\theta_i$
    idx_theta = torch.outer(seq_idx, theta)

    # Compute cache. Because polar only takes float32 or float64, we need to cast
    # when working with 16 bit floats (float16 or bfloat16)
    dtypes_requiring_casting = [torch.float16, torch.bfloat16, torch.int8]
    working_dtype = (
        torch.float32 if dtype in dtypes_requiring_casting else dtype
    )
    complex_dtype = (
        torch.complex32 if dtype in dtypes_requiring_casting else torch.complex64
    )
    cache = torch.polar(
        torch.ones_like(idx_theta).to(working_dtype), idx_theta.to(working_dtype)
    ).to(complex_dtype)
    return cache


def apply_rope(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
    x = x.transpose(1, 2)

    # truncate to support variable sizes
    T = x.size(1)
    rope_cache = rope_cache[:T]
    # cast because `view_as_complex` does not support 16 bit tensors
    xc = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
    rope_cache = rope_cache.view(1, xc.size(1), 1, xc.size(3))
    x_out = torch.view_as_real(xc * rope_cache).flatten(3)
    return x_out.transpose(1, 2).type_as(x)