weaviate
107 строк · 2.5 Кб
1// _ _
2// __ _____ __ ___ ___ __ _| |_ ___
3// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
4// \ V V / __/ (_| |\ V /| | (_| | || __/
5// \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
6//
7// Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
8//
9// CONTACT: hello@weaviate.io
10//
11
12package vectorizer13
14import (15"context"16
17"github.com/weaviate/weaviate/modules/text2vec-huggingface/ent"18)
19
20type fakeClient struct {21lastInput string22lastConfig ent.VectorizationConfig23}
24
25func (c *fakeClient) Vectorize(ctx context.Context,26text string, cfg ent.VectorizationConfig,27) (*ent.VectorizationResult, error) {28c.lastInput = text29c.lastConfig = cfg30return &ent.VectorizationResult{31Vector: []float32{0, 1, 2, 3},32Dimensions: 4,33Text: text,34}, nil35}
36
37func (c *fakeClient) VectorizeQuery(ctx context.Context,38text string, cfg ent.VectorizationConfig,39) (*ent.VectorizationResult, error) {40c.lastInput = text41c.lastConfig = cfg42return &ent.VectorizationResult{43Vector: []float32{0.1, 1.1, 2.1, 3.1},44Dimensions: 4,45Text: text,46}, nil47}
48
49type fakeClassConfig struct {50classConfig map[string]interface{}51vectorizeClassName bool52vectorizePropertyName bool53skippedProperty string54excludedProperty string55// module specific settings56model, passageModel, queryModel string57waitForModel, useGPU, useCache bool58endpointURL string59}
60
61func (f fakeClassConfig) Class() map[string]interface{} {62if len(f.classConfig) > 0 {63return f.classConfig64}65classSettings := map[string]interface{}{66"vectorizeClassName": f.vectorizeClassName,67"model": f.model,68"passageModel": f.passageModel,69"queryModel": f.queryModel,70"waitForModel": f.waitForModel,71"useGPU": f.useGPU,72"useCache": f.useCache,73"endpointURL": f.endpointURL,74}75return classSettings76}
77
78func (f fakeClassConfig) ClassByModuleName(moduleName string) map[string]interface{} {79return f.classConfig80}
81
82func (f fakeClassConfig) Property(propName string) map[string]interface{} {83if propName == f.skippedProperty {84return map[string]interface{}{85"skip": true,86}87}88if propName == f.excludedProperty {89return map[string]interface{}{90"vectorizePropertyName": false,91}92}93if f.vectorizePropertyName {94return map[string]interface{}{95"vectorizePropertyName": true,96}97}98return nil99}
100
101func (f fakeClassConfig) Tenant() string {102return ""103}
104
105func (f fakeClassConfig) TargetVector() string {106return ""107}
108