def extract_features(
        self,
        source: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        mask: bool = False,
        ret_conv: bool = False,
        output_layer: Optional[int] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Extract features for only speech input"""
        res = self.forward(
            source,
            padding_mask=padding_mask,
            mask=mask,
            features_only=True,
            output_layer=output_layer,
        )
        x = res["x"] # B x T x D
        padding_mask = res["padding_mask"]

        if self.add_unit_encoder:
            src_tokens, x, _ = self.convert_embeddings(
                x,
                padding_mask,
                mix_with_unit=False,
                use_pred_unit=False,
            )
            encoder_out = self.unit_encoder(
                src_tokens,
                token_embeddings=x,
                return_all_hiddens=output_layer is not None
            )
            res["x"] = encoder_out['encoder_out'][0].transpose(0, 1)  # (B, T, D)
        
        feature = res["features"] if ret_conv else res["x"]
        if output_layer is not None:
            feature = encoder_out['encoder_states']

        return feature, padding_mask
Extract features for only speech input
speechlm/speechlm/models/speechlm.py