correct labels

Author: ngxson

convert_hf_to_gguf.py

@@ -3095,7 +3095,7 @@ def set_gguf_parameters(self):
         if is_rerank:
             self.gguf_writer.add_pooling_type(gguf.PoolingType.RANK)
             self.gguf_writer.add_sep_token_id(self.sep_token_id)
-            self.gguf_writer.add_uint32(gguf.Keys.Classifier.OUTPUT_LABELS, 2)
+            self.gguf_writer.add_classifier_output_labels(["yes", "no"])
 
     def _get_cls_out_tensor(self, data_torch: Tensor) -> Tensor:
         # extract "yes" and "no" tokens from the output lm_head tensor

src/llama-model.cpp

@@ -2470,7 +2470,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
                     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 
                     // output rerank
-                    cls_out = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), {n_embd, 2}, TENSOR_NOT_REQUIRED);
+                    cls_out = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), {n_embd, hparams.n_cls_out}, TENSOR_NOT_REQUIRED);
 
                     // output
                     output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);