weaviate

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

12
package autocorrect
13

14
import (
15
	"context"
16
	"strings"
17

18
	"github.com/weaviate/weaviate/modules/text-spellcheck/ent"
19
)
20

21
type spellCheckClient interface {
22
	Check(ctx context.Context, text []string) (*ent.SpellCheckResult, error)
23
}
24

25
type AutocorrectTransformer struct {
26
	spellCheckClient spellCheckClient
27
}
28

29
func New(spellCheckClient spellCheckClient) *AutocorrectTransformer {
30
	return &AutocorrectTransformer{spellCheckClient}
31
}
32

33
func (t *AutocorrectTransformer) Transform(in []string) ([]string, error) {
34
	spellCheckResult, err := t.spellCheckClient.Check(context.Background(), in)
35
	if err != nil {
36
		return nil, err
37
	}
38
	result := make([]string, len(in))
39
	changes := spellCheckResult.Changes
40
	for i, txt := range spellCheckResult.Text {
41
		didYouMean := txt
42
		for _, change := range changes {
43
			didYouMean = strings.ReplaceAll(strings.ToLower(didYouMean), change.Original, change.Correction)
44
		}
45
		result[i] = didYouMean
46
	}
47
	return result, nil
48
}
49