From 63861f58cc6bad7d512825badd5154ccd7b32826 Mon Sep 17 00:00:00 2001
From: Jeffrey Morgan <jmorganca@gmail.com>
Date: Tue, 20 Feb 2024 21:37:29 -0500
Subject: [PATCH] Support for `bert` and `nomic-bert` embedding models

---
 server/images.go |  4 ++++
 server/routes.go | 10 ++++++++++
 2 files changed, 14 insertions(+)

diff --git a/server/images.go b/server/images.go
index a000c5c5..83ef7920 100644
--- a/server/images.go
+++ b/server/images.go
@@ -52,6 +52,10 @@ type Model struct {
 	Messages       []Message
 }
 
+func (m *Model) IsEmbedding() bool {
+	return slices.Contains(m.Config.ModelFamilies, "bert") || slices.Contains(m.Config.ModelFamilies, "nomic-bert")
+}
+
 type Message struct {
 	Role    string `json:"role"`
 	Content string `json:"content"`
diff --git a/server/routes.go b/server/routes.go
index c769a876..254d26ad 100644
--- a/server/routes.go
+++ b/server/routes.go
@@ -191,6 +191,11 @@ func GenerateHandler(c *gin.Context) {
 		return
 	}
 
+	if model.IsEmbedding() {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model does not support generate"})
+		return
+	}
+
 	opts, err := modelOptions(model, req.Options)
 	if err != nil {
 		if errors.Is(err, api.ErrInvalidOpts) {
@@ -1143,6 +1148,11 @@ func ChatHandler(c *gin.Context) {
 		return
 	}
 
+	if model.IsEmbedding() {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model does not support chat"})
+		return
+	}
+
 	opts, err := modelOptions(model, req.Options)
 	if err != nil {
 		if errors.Is(err, api.ErrInvalidOpts) {