graph : remove model reference from build_pooling

ggml-ci

Author: ggerganov

src/llama-graph.cpp

@@ -1992,7 +1992,12 @@ ggml_tensor * llm_graph_context::build_rwkv_channel_mix(
     return cur;
 }
 
-void llm_graph_context::build_pooling(ggml_cgraph * gf) const {
+void llm_graph_context::build_pooling(
+        ggml_cgraph * gf,
+        ggml_tensor * cls,
+        ggml_tensor * cls_b,
+        ggml_tensor * cls_out,
+        ggml_tensor * cls_out_b) const {
     if (!cparams.embeddings) {
         return;
     }
@@ -2036,18 +2041,18 @@ void llm_graph_context::build_pooling(ggml_cgraph * gf) const {
 
                 // classification head
                 // https://github.com/huggingface/transformers/blob/5af7d41e49bbfc8319f462eb45253dcb3863dfb7/src/transformers/models/roberta/modeling_roberta.py#L1566
-                GGML_ASSERT(model.cls   != nullptr);
-                GGML_ASSERT(model.cls_b != nullptr);
+                GGML_ASSERT(cls   != nullptr);
+                GGML_ASSERT(cls_b != nullptr);
 
-                cur = ggml_add (ctx0, ggml_mul_mat(ctx0, model.cls, inp), model.cls_b);
+                cur = ggml_add (ctx0, ggml_mul_mat(ctx0, cls, inp), cls_b);
                 cur = ggml_tanh(ctx0, cur);
 
                 // some models don't have `cls_out`, for example: https://huggingface.co/jinaai/jina-reranker-v1-tiny-en
                 // https://huggingface.co/jinaai/jina-reranker-v1-tiny-en/blob/cb5347e43979c3084a890e3f99491952603ae1b7/modeling_bert.py#L884-L896
-                if (model.cls_out) {
-                    GGML_ASSERT(model.cls_out_b != nullptr);
+                if (cls_out) {
+                    GGML_ASSERT(cls_out_b != nullptr);
 
-                    cur = ggml_add (ctx0, ggml_mul_mat(ctx0, model.cls_out, cur), model.cls_out_b);
+                    cur = ggml_add (ctx0, ggml_mul_mat(ctx0, cls_out, cur), cls_out_b);
                 }
             } break;
         default:

src/llama-graph.h

@@ -617,5 +617,10 @@ struct llm_graph_context {
     // pooling
     //
 
-    void build_pooling(ggml_cgraph * gf) const;
+    void build_pooling(
+            ggml_cgraph * gf,
+            ggml_tensor * cls,
+            ggml_tensor * cls_b,
+            ggml_tensor * cls_out,
+            ggml_tensor * cls_out_b) const;
 };

src/llama-model.cpp

@@ -10583,7 +10583,7 @@ llm_graph_result_ptr llama_model::build_graph(
     }
 
     // add on pooling layer
-    llm->build_pooling(gf);
+    llm->build_pooling(gf, cls, cls_b, cls_out, cls_out_b);
 
     return std::move(llm->res);
 }