def embed_tokens()

in summarize_from_feedback/models/transformer.py [0:0]

• 32 lines of code
• 5 McCabe index (conditional complexity)

    def embed_tokens(self, x, ctx_len=0, act_dtype=torch.float16):
        tokens = x
        assert isinstance(
            tokens, (torch.LongTensor, torch.cuda.LongTensor)
        ), f"Tokens should be type long: {getattr(tokens, 'dtype', type(tokens))}"
        assert (0 <= tokens).all() and (
            tokens < self.n_vocab
        ).all(), f"{tokens.max()} >= {self.n_vocab} or {tokens.min()} < 0"

        emb = F.embedding(
            tokens,
            self.embedding.weight,
            self.embedding.padding_idx,
            self.embedding.max_norm,
            self.embedding.norm_type,
            self.embedding.scale_grad_by_freq,
            self.embedding.sparse,
        ).type(act_dtype)

        if self.include_pos_embeddings:
            pos_emb = self.position_embedding(
                torch.tensor(ctx_len), torch.tensor(ctx_len + tokens.size(-1))
            ).type(act_dtype)

            embedded_tokens = self.embed_dropout(emb) + self.pos_emb_dropout(pos_emb)
        else:
            embedded_tokens = self.embed_dropout(emb)

        if self.mp_comm is not None:
            comm = self.mp_comm
            if len(comm.ranks) == 1:
                result = [embedded_tokens]
            else:
                result = comm.all_gather_no_backward(embedded_tokens, "input_gather")
            embedded_tokens = torch.cat(result, dim=-1)
        return embedded_tokens