github.com/hupe1980/go-huggingface@v0.0.15/sentence_similarity.go (about)

     1  package huggingface
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"errors"
     7  )
     8  
     9  // SentenceSimilarityInputs represents the inputs for sentence similarity computation.
    10  type SentenceSimilarityInputs struct {
    11  	SourceSentence string   `json:"source_sentence"`
    12  	Sentences      []string `json:"sentences"`
    13  }
    14  
    15  // SentenceSimilarityRequest represents a request for sentence similarity computation.
    16  type SentenceSimilarityRequest struct {
    17  	Inputs  SentenceSimilarityInputs `json:"inputs"`
    18  	Options Options                  `json:"options,omitempty"`
    19  	Model   string                   `json:"-"`
    20  }
    21  
    22  // SentenceSimilarityResponse represents the response for a sentence similarity computation request.
    23  type SentenceSimilarityResponse []float32
    24  
    25  // SentenceSimilarity sends a sentence similarity computation request to the InferenceClient
    26  // and returns the sentence similarity response.
    27  func (ic *InferenceClient) SentenceSimilarity(ctx context.Context, req *SentenceSimilarityRequest) (SentenceSimilarityResponse, error) {
    28  	if len(req.Inputs.SourceSentence) == 0 || len(req.Inputs.Sentences) == 0 {
    29  		return nil, errors.New("sourceSentence and sentences are required")
    30  	}
    31  
    32  	body, err := ic.post(ctx, req.Model, "sentence-similarity", req)
    33  
    34  	if err != nil {
    35  		return nil, err
    36  	}
    37  
    38  	sentenceSimilarityResponse := SentenceSimilarityResponse{}
    39  
    40  	if err := json.Unmarshal(body, &sentenceSimilarityResponse); err != nil {
    41  		return nil, err
    42  	}
    43  
    44  	return sentenceSimilarityResponse, nil
    45  }