From 4524290e87b8e107cc2b56e1251751546f4b9051 Mon Sep 17 00:00:00 2001
From: Douglas Hanley <thesecretaryofwar@gmail.com>
Date: Thu, 15 Feb 2024 11:21:49 -0600
Subject: Use correct type of pooling for embedding models (#5500)

Use correct type of pooling for embedding models
---
 gguf-py/gguf/constants.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

(limited to 'gguf-py/gguf/constants.py')

diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index 9986ce9d..114a9a97 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -40,7 +40,7 @@ class Keys:
         TENSOR_DATA_LAYOUT    = "{arch}.tensor_data_layout"
         EXPERT_COUNT          = "{arch}.expert_count"
         EXPERT_USED_COUNT     = "{arch}.expert_used_count"
-        POOLING_LAYER         = "{arch}.pooling_layer"
+        POOLING_TYPE          = "{arch}.pooling_type"
 
     class Attention:
         HEAD_COUNT        = "{arch}.attention.head_count"
@@ -561,6 +561,12 @@ class RopeScalingType(Enum):
     YARN   = 'yarn'
 
 
+class PoolingType(IntEnum):
+    NONE = 0
+    MEAN = 1
+    CLS  = 2
+
+
 class GGMLQuantizationType(IntEnum):
     F32  = 0
     F16  = 1
-- 
cgit v1.2.3