From da8e2a04479f96ad9c57eaf25ed26b79b239b05c Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Fri, 14 Jun 2024 14:57:49 -0700
Subject: [PATCH] use kvs to detect embedding models

---
 server/images.go | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/server/images.go b/server/images.go
index 5cd0a7a53..a62991f16 100644
--- a/server/images.go
+++ b/server/images.go
@@ -66,7 +66,21 @@ func (m *Model) Has(caps ...Capability) bool {
 	for _, cap := range caps {
 		switch cap {
 		case CapabilityCompletion:
-			if slices.Contains(m.Config.ModelFamilies, "bert") || slices.Contains(m.Config.ModelFamilies, "nomic-bert") {
+			f, err := os.Open(m.ModelPath)
+			if err != nil {
+				slog.Error("couldn't open model file", "error", err)
+				continue
+			}
+			defer f.Close()
+
+			// TODO(mxyng): decode the GGML into model to avoid doing this multiple times
+			ggml, _, err := llm.DecodeGGML(f, 0)
+			if err != nil {
+				slog.Error("couldn't decode ggml", "error", err)
+				continue
+			}
+
+			if _, ok := ggml.KV()[fmt.Sprintf("%s.pooling_type", ggml.KV().Architecture())]; ok {
 				return false
 			}
 		default: