From 95845d17ec61267ba902d285715c17b168ffd80e Mon Sep 17 00:00:00 2001 From: Jared Van Bortel Date: Fri, 1 Mar 2024 15:19:59 -0500 Subject: [PATCH] convert-hf : make actual types match annotations --- convert-hf-to-gguf.py | 2 +- gguf-py/gguf/gguf_writer.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/convert-hf-to-gguf.py b/convert-hf-to-gguf.py index 28b92ac38..470d494e3 100755 --- a/convert-hf-to-gguf.py +++ b/convert-hf-to-gguf.py @@ -1679,7 +1679,7 @@ class BertModel(Model): else: raise NotImplementedError("Only MEAN and CLS pooling types supported") - self.gguf_writer.add_pooling_type(pooling_type.value) + self.gguf_writer.add_pooling_type(pooling_type) def set_vocab(self): path = self.dir_model diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index e4681475c..801160832 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -362,7 +362,7 @@ class GGUFWriter: self.add_bool(Keys.Attention.CAUSAL.format(arch=self.arch), value) def add_pooling_type(self, value: PoolingType) -> None: - self.add_uint32(Keys.LLM.POOLING_TYPE.format(arch=self.arch), value) + self.add_uint32(Keys.LLM.POOLING_TYPE.format(arch=self.arch), value.value) def add_rope_dimension_count(self, count: int) -> None: self.add_uint32(Keys.Rope.DIMENSION_COUNT.format(arch=self.arch), count)