Class MetaModel

model/meta.py:12–175

class MetaModel(nn.Module):

    def __init__(self, llama_type, llama_config, llama_ckpt_dir=None, tokenizer_path=None):
        super().__init__()

        self.criterion = torch.nn.CrossEntropyLoss(ignore_index=0)

        ModelArgs = LLM.__dict__[llama_type].ModelArgs
        Transformer = LLM.__dict__[llama_type].Transformer

        with open(llama_config, "r") as f:
            params = json.loads(f.read())
        model_args: ModelArgs = ModelArgs(
            max_seq_len=2048, max_batch_size=32, **params
        )
        self.tokenizer = Tokenizer(model_path=tokenizer_path)
        model_args.vocab_size = self.tokenizer.n_words

        model = Transformer(model_args)
        mp_rank = fs_init.get_model_parallel_rank()
        if llama_ckpt_dir is not None:
            ckpt_path = os.path.join(llama_ckpt_dir, f"consolidated.{mp_rank:02d}.pth")
            if os.path.exists(ckpt_path):
                checkpoint = torch.load(ckpt_path, map_location="cpu")
                msg = model.load_state_dict(checkpoint, strict=False)
                print(msg)
            else:
                print(f'Checkpoint not found at {ckpt_path}')
        self.llma = model
        for name, param in self.named_parameters():
            if param.requires_grad:
               print(f"Trainable param: {name}, {param.shape}, {param.dtype}")
        count = sum(p.numel() for p in self.parameters() if p.requires_grad)
        print(f"Parameter count : {count}")

    def forward(self, examples, labels, image=None, modal='image'):
        output = self.llma(examples, image=image, modal=modal)
        output = output[:, :-1, :]
        labels = labels[:, 1:]

        if labels.sum() == 0:
            c_loss = output.mean() * 0
        else:
            c_loss = self.criterion(output.reshape(-1, 32000), labels.flatten())

        return c_loss

    def generate(
        self,
        prompts: List[str],
        images,
        max_gen_len: int,
        temperature: float = 0.8,
        top_p: float = 0.95,
        modal = ['image'],
    ) -> List[str]:
        bsz = len(prompts)
        params = self.llma.params