Merge branch 'feat/averiewang/support-new-model' into 'main' (merge request !61)

feat: support new embedding model

Author: rogersqsliu

CHANGELOG.md

@@ -1,5 +1,7 @@
 # Changelog
 
+## v1.4.6
+* 支持BAAI/bge-m3新模型，也支持创建embedding collection时使用string直接设置模型
 
 ## v1.4.5
 * 更换依赖cgo的分词包，为纯go实现的分词包，以更好的支持跨平台编译

example/embedding_demo/main.go

@@ -109,7 +109,7 @@ func (d *EmbeddingDemo) CreateDBAndCollection(ctx context.Context, database, col
 	index.FilterIndex = append(index.FilterIndex, tcvectordb.FilterIndex{FieldName: "bookName", FieldType: tcvectordb.String, IndexType: tcvectordb.FILTER})
 	index.FilterIndex = append(index.FilterIndex, tcvectordb.FilterIndex{FieldName: "page", FieldType: tcvectordb.Uint64, IndexType: tcvectordb.FILTER})
 
-	ebd := &tcvectordb.Embedding{VectorField: "vector", Field: "text", Model: tcvectordb.BGE_BASE_ZH}
+	ebd := &tcvectordb.Embedding{VectorField: "vector", Field: "text", ModelName: "bge-base-zh"}
 	// 第二步：创建 Collection
 	// 创建支持 Embedding 的 Collection
 	db.WithTimeout(time.Second * 30)

tcvectordb/base_collection.go

@@ -145,6 +145,9 @@ func (i *implementerCollection) CreateCollection(ctx context.Context, name strin
 			req.Embedding.Field = param.Embedding.Field
 			req.Embedding.VectorField = param.Embedding.VectorField
 			req.Embedding.Model = string(param.Embedding.Model)
+			if param.Embedding.ModelName != "" {
+				req.Embedding.Model = param.Embedding.ModelName
+			}
 		}
 		if param.TtlConfig != nil {
 			req.TtlConfig = new(collection.TtlConfig)
@@ -315,6 +318,7 @@ func (i *implementerCollection) toCollection(collectionItem *collection.Describe
 		coll.Embedding.Field = collectionItem.Embedding.Field
 		coll.Embedding.VectorField = collectionItem.Embedding.VectorField
 		coll.Embedding.Model = EmbeddingModel(collectionItem.Embedding.Model)
+		coll.Embedding.ModelName = collectionItem.Embedding.Model
 		coll.Embedding.Enabled = collectionItem.Embedding.Status == "enabled"
 	}
 	if collectionItem.TtlConfig != nil {
@@ -459,10 +463,12 @@ func (c *Collection) WithTimeout(t time.Duration) {
 }
 
 type Embedding struct {
-	Field       string         `json:"field,omitempty"`
-	VectorField string         `json:"vectorField,omitempty"`
-	Model       EmbeddingModel `json:"model,omitempty"`
-	Enabled     bool           `json:"enabled,omitempty"` // 返回数据
+	Field       string `json:"field,omitempty"`
+	VectorField string `json:"vectorField,omitempty"`
+	// Deprecated: Use ModelName instead.
+	Model     EmbeddingModel `json:"model,omitempty"`
+	ModelName string         `json:"modelName,omitempty"`
+	Enabled   bool           `json:"enabled,omitempty"` // 返回数据
 }
 
 type IndexStatus struct {

tcvectordb/consts.go

@@ -71,6 +71,8 @@ const (
 	E5_LARGE_V2 EmbeddingModel = "e5-large-v2"
 	// TEXT2VEC_LARGE_CHINESE 1024
 	TEXT2VEC_LARGE_CHINESE EmbeddingModel = "text2vec-large-chinese"
+	// BAAI_BGE_M3 1024
+	BAAI_BGE_M3 EmbeddingModel = "BAAI/bge-m3"
 )
 
 type ReadConsistency string

tcvectordb/rpc_base_collection.go

@@ -101,6 +101,9 @@ func (r *rpcImplementerCollection) CreateCollection(ctx context.Context, name st
 				VectorField: param.Embedding.VectorField,
 				ModelName:   string(param.Embedding.Model),
 			}
+			if param.Embedding.ModelName != "" {
+				req.EmbeddingParams.ModelName = param.Embedding.ModelName
+			}
 		}
 		if param.TtlConfig != nil {
 			req.TtlConfig = &olama.TTLConfig{
@@ -246,6 +249,7 @@ func (r *rpcImplementerCollection) toCollection(collectionItem *olama.CreateColl
 		coll.Embedding.Field = collectionItem.EmbeddingParams.Field
 		coll.Embedding.VectorField = collectionItem.EmbeddingParams.VectorField
 		coll.Embedding.Model = EmbeddingModel(collectionItem.EmbeddingParams.ModelName)
+		coll.Embedding.ModelName = collectionItem.EmbeddingParams.ModelName
 		coll.Embedding.Enabled = false
 	}
 	if collectionItem.TtlConfig != nil {

tcvectordb/version.go

@@ -18,4 +18,4 @@
 
 package tcvectordb
 
-const SDKVersion = "v1.4.5"
+const SDKVersion = "v1.4.6"

test/embedding_test.go

@@ -37,7 +37,7 @@ func TestCreateCollectionWithEmbedding(t *testing.T) {
 		Embedding: &tcvectordb.Embedding{
 			Field:       "segment",
 			VectorField: "vector",
-			Model:       tcvectordb.BGE_BASE_ZH,
+			ModelName:   "BAAI/bge-m3",
 		},
 	}
 

test/normal_test.go

@@ -426,3 +426,8 @@ func TestJson(t *testing.T) {
 	temp1 := temp["shardNum"]
 	println(fmt.Sprintf("%T, %v", temp1, temp1))
 }
+
+func TestEmeb(t *testing.T) {
+	model := "model_bge"
+	println(tcvectordb.EmbeddingModel(model))
+}