Pārlūkot izejas kodu

use kvs to detect embedding models

Michael Yang 10 mēneši atpakaļ
vecāks
revīzija
da8e2a0447
1 mainītis faili ar 15 papildinājumiem un 1 dzēšanām
  1. 15 1
      server/images.go

+ 15 - 1
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: