github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/client/reranker.go (about)

     1  package client
     2  
     3  import (
     4  	"encoding/json"
     5  
     6  	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
     7  )
     8  
     9  const (
    10  	rerankType    = "strategy"
    11  	rerankParams  = "params"
    12  	rffParam      = "k"
    13  	weightedParam = "weights"
    14  
    15  	rrfRerankType      = `rrf`
    16  	weightedRerankType = `weighted`
    17  )
    18  
    19  type Reranker interface {
    20  	GetParams() []*commonpb.KeyValuePair
    21  }
    22  
    23  type rrfReranker struct {
    24  	K float64 `json:"k,omitempty"`
    25  }
    26  
    27  func (r *rrfReranker) WithK(k float64) *rrfReranker {
    28  	r.K = k
    29  	return r
    30  }
    31  
    32  func (r *rrfReranker) GetParams() []*commonpb.KeyValuePair {
    33  	bs, _ := json.Marshal(r)
    34  
    35  	return []*commonpb.KeyValuePair{
    36  		{Key: rerankType, Value: rrfRerankType},
    37  		{Key: rerankParams, Value: string(bs)},
    38  	}
    39  }
    40  
    41  func NewRRFReranker() *rrfReranker {
    42  	return &rrfReranker{K: 60}
    43  }
    44  
    45  type weightedReranker struct {
    46  	Weights []float64 `json:"weights,omitempty"`
    47  }
    48  
    49  func (r *weightedReranker) GetParams() []*commonpb.KeyValuePair {
    50  	bs, _ := json.Marshal(r)
    51  
    52  	return []*commonpb.KeyValuePair{
    53  		{Key: rerankType, Value: weightedRerankType},
    54  		{Key: rerankParams, Value: string(bs)},
    55  	}
    56  }
    57  
    58  func NewWeightedReranker(weights []float64) *weightedReranker {
    59  	return &weightedReranker{
    60  		Weights: weights,
    61  	}
    62  }