KwangHwi/quantization_v2

1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.2# SPDX-License-Identifier: Apache-2.0.3from torch import Tensor4import torch5 6 7def pool(last_hidden_states: Tensor, attention_mask: Tensor, pool_type: str) -> Tensor:8    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)9 10    if pool_type == "avg":11        emb = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]12    elif pool_type == "weighted_avg":13        emb = last_hidden.sum(dim=1)14    elif pool_type == "cls":15        emb = last_hidden[:, 0]16    elif pool_type == "last":17        left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]18        if left_padding:19            emb = last_hidden[:, -1]20        else:21            sequence_lengths = attention_mask.sum(dim=1) - 122            batch_size = last_hidden.shape[0]23            emb = last_hidden[24                torch.arange(batch_size, device=last_hidden.device), sequence_lengths25            ]26    else:27        raise ValueError(f"pool_type {pool_type} not supported")28 29    return emb30