Support for bert and nomic-bert embedding models

This commit is contained in:
Jeffrey Morgan 2024-02-20 21:37:29 -05:00 committed by GitHub
parent f0425d3de9
commit 63861f58cc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 14 additions and 0 deletions

View file

@ -52,6 +52,10 @@ type Model struct {
Messages []Message Messages []Message
} }
func (m *Model) IsEmbedding() bool {
return slices.Contains(m.Config.ModelFamilies, "bert") || slices.Contains(m.Config.ModelFamilies, "nomic-bert")
}
type Message struct { type Message struct {
Role string `json:"role"` Role string `json:"role"`
Content string `json:"content"` Content string `json:"content"`

View file

@ -191,6 +191,11 @@ func GenerateHandler(c *gin.Context) {
return return
} }
if model.IsEmbedding() {
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model does not support generate"})
return
}
opts, err := modelOptions(model, req.Options) opts, err := modelOptions(model, req.Options)
if err != nil { if err != nil {
if errors.Is(err, api.ErrInvalidOpts) { if errors.Is(err, api.ErrInvalidOpts) {
@ -1143,6 +1148,11 @@ func ChatHandler(c *gin.Context) {
return return
} }
if model.IsEmbedding() {
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model does not support chat"})
return
}
opts, err := modelOptions(model, req.Options) opts, err := modelOptions(model, req.Options)
if err != nil { if err != nil {
if errors.Is(err, api.ErrInvalidOpts) { if errors.Is(err, api.ErrInvalidOpts) {