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

     1  package huggingface
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"errors"
     7  )
     8  
     9  type SummarizationParameters struct {
    10  	// (Default: None). Integer to define the minimum length in tokens of the output summary.
    11  	MinLength *int `json:"min_length,omitempty"`
    12  
    13  	// (Default: None). Integer to define the maximum length in tokens of the output summary.
    14  	MaxLength *int `json:"max_length,omitempty"`
    15  
    16  	// (Default: None). Integer to define the top tokens considered within the sample operation to create
    17  	// new text.
    18  	TopK *int `json:"top_k,omitempty"`
    19  
    20  	// (Default: None). Float to define the tokens that are within the sample` operation of text generation.
    21  	// Add tokens in the sample for more probable to least probable until the sum of the probabilities is
    22  	// greater than top_p.
    23  	TopP *float64 `json:"top_p,omitempty"`
    24  
    25  	// (Default: 1.0). Float (0.0-100.0). The temperature of the sampling operation. 1 means regular sampling,
    26  	// 0 mens top_k=1, 100.0 is getting closer to uniform probability.
    27  	Temperature *float64 `json:"temperature,omitempty"`
    28  
    29  	// (Default: None). Float (0.0-100.0). The more a token is used within generation the more it is penalized
    30  	// to not be picked in successive generation passes.
    31  	RepetitionPenalty *float64 `json:"repetitionpenalty,omitempty"`
    32  
    33  	// (Default: None). Float (0-120.0). The amount of time in seconds that the query should take maximum.
    34  	// Network can cause some overhead so it will be a soft limit.
    35  	MaxTime *float64 `json:"maxtime,omitempty"`
    36  }
    37  
    38  type SummarizationRequest struct {
    39  	// String to be summarized
    40  	Inputs     []string                `json:"inputs"`
    41  	Parameters SummarizationParameters `json:"parameters,omitempty"`
    42  	Options    Options                 `json:"options,omitempty"`
    43  	Model      string                  `json:"-"`
    44  }
    45  
    46  type SummarizationResponse []struct {
    47  	// The summarized input string
    48  	SummaryText string `json:"summary_text,omitempty"`
    49  }
    50  
    51  // Summarization performs text summarization using the specified model.
    52  // It sends a POST request to the Hugging Face inference endpoint with the provided inputs.
    53  // The response contains the generated summary or an error if the request fails.
    54  func (ic *InferenceClient) Summarization(ctx context.Context, req *SummarizationRequest) (SummarizationResponse, error) {
    55  	if len(req.Inputs) == 0 {
    56  		return nil, errors.New("inputs are required")
    57  	}
    58  
    59  	body, err := ic.post(ctx, req.Model, "summarization", req)
    60  	if err != nil {
    61  		return nil, err
    62  	}
    63  
    64  	summarizationResponse := SummarizationResponse{}
    65  	if err := json.Unmarshal(body, &summarizationResponse); err != nil {
    66  		return nil, err
    67  	}
    68  
    69  	return summarizationResponse, nil
    70  }