weaviate
60 строк · 1.6 Кб
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 vectorizer
13
14import (
15"context"
16
17"github.com/weaviate/weaviate/entities/models"
18"github.com/weaviate/weaviate/entities/moduletools"
19"github.com/weaviate/weaviate/modules/text2vec-gpt4all/ent"
20objectsvectorizer "github.com/weaviate/weaviate/usecases/modulecomponents/vectorizer"
21)
22
23type Vectorizer struct {
24client Client
25objectVectorizer *objectsvectorizer.ObjectVectorizer
26}
27
28func New(client Client) *Vectorizer {
29return &Vectorizer{
30client: client,
31objectVectorizer: objectsvectorizer.New(),
32}
33}
34
35type Client interface {
36Vectorize(ctx context.Context, text string) (*ent.VectorizationResult, error)
37}
38
39type ClassSettings interface {
40PropertyIndexed(property string) bool
41VectorizeClassName() bool
42VectorizePropertyName(propertyName string) bool
43}
44
45func (v *Vectorizer) Object(ctx context.Context, object *models.Object, cfg moduletools.ClassConfig,
46) ([]float32, models.AdditionalProperties, error) {
47vec, err := v.object(ctx, object, cfg)
48return vec, nil, err
49}
50
51func (v *Vectorizer) object(ctx context.Context, object *models.Object, cfg moduletools.ClassConfig,
52) ([]float32, error) {
53text := v.objectVectorizer.Texts(ctx, object, NewClassSettings(cfg))
54res, err := v.client.Vectorize(ctx, text)
55if err != nil {
56return nil, err
57}
58
59return res.Vector, nil
60}
61