Skip to content

Commit

Permalink
Makes PoolingConfig a dataclass
Browse files Browse the repository at this point in the history
Signed-off-by: Flavia Beo <flavia.beo@ibm.com>
  • Loading branch information
flaviabeo committed Oct 22, 2024
1 parent 2c5653c commit 55779a3
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions vllm/model_executor/layers/pooler.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from enum import IntEnum
from dataclasses import dataclass

import torch
import torch.nn as nn
Expand All @@ -15,9 +16,11 @@ class PoolingType(IntEnum):
CLS = 2
MEAN = 3


class PoolingConfig():
"""A class that configures the pooling operation.
@dataclass
class PoolingConfig:
"""A class that configures the pooling operation which
only applies to sentence-transformers models.
More at: https://www.sbert.net/
Attributes:
pooling_type (str): The type of pooling to use.
Expand Down

0 comments on commit 55779a3

Please sign in to comment.