dgraph/dgraphapi/vector.go at 4198b0fac88235bf31e855f303971bfbdddc185f · dgraph-io/dgraph · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
/*
 * SPDX-FileCopyrightText: © Hypermode Inc. <[email protected]>
 * SPDX-License-Identifier: Apache-2.0
 */

package dgraphapi

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"strings"

	"github.com/dgraph-io/dgo/v250/protos/api"
)

func GenerateRandomVector(size int) []float32 {
	vector := make([]float32, size)
	for i := range size {
		vector[i] = rand.Float32() * 10
	}
	return vector
}

func formatVector(label string, vector []float32, index int) string {
	vectorString := fmt.Sprintf(`"[%s]"`, strings.Trim(strings.Join(strings.Fields(fmt.Sprint(vector)), ", "), "[]"))
	return fmt.Sprintf("<0x%x> <%s> %s . \n", index+10, label, vectorString)
}

func GenerateRandomVectors(lowerLimit, uppermLimit, vectorSize int, label string) (string, [][]float32) {
	var builder strings.Builder
	var vectors [][]float32
	// builder.WriteString("`")
	for i := lowerLimit; i < uppermLimit; i++ {
		randomVector := GenerateRandomVector(vectorSize)
		vectors = append(vectors, randomVector)
		formattedVector := formatVector(label, randomVector, i)
		builder.WriteString(formattedVector)
	}

	return builder.String(), vectors
}

func (gc *GrpcClient) QueryMultipleVectorsUsingSimilarTo(vector []float32, pred string, topK int) ([][]float32, error) {
	vectorQuery := fmt.Sprintf(`
	 {
		 vector(func: similar_to(%v, %v, "%v")) {
				uid
				%v
		  }
	 }`, pred, topK, vector, pred)
	resp, err := gc.Query(vectorQuery)

	if err != nil {
		return [][]float32{}, err
	}

	return UnmarshalVectorResp(resp)
}

func (gc *GrpcClient) QuerySingleVectorsUsingUid(uid, pred string) ([][]float32, error) {
	vectorQuery := fmt.Sprintf(`
	    {
		 vector(func: uid(%v)) {
			uid
			%v
		     }
	    }`, uid[1:len(uid)-1], pred)

	resp, err := gc.Query(vectorQuery)
	if err != nil {
		return [][]float32{}, err
	}

	return UnmarshalVectorResp(resp)
}

func UnmarshalVectorResp(resp *api.Response) ([][]float32, error) {
	type Data struct {
		Vector []struct {
			UID                 string    `json:"uid"`
			ProjectDescriptionV []float32 `json:"project_description_v"`
		} `json:"vector"`
	}
	var data Data
	if err := json.Unmarshal(resp.Json, &data); err != nil {
		return nil, err
	}

	var vectors [][]float32
	for _, item := range data.Vector {
		vectors = append(vectors, item.ProjectDescriptionV)
	}
	return vectors, nil
}