File size: 1,497 Bytes
4674012 | 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 | package dto
import (
"fmt"
"strings"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
)
type RerankRequest struct {
Documents []any `json:"documents"`
Query string `json:"query"`
Model string `json:"model"`
TopN int `json:"top_n,omitempty"`
ReturnDocuments *bool `json:"return_documents,omitempty"`
MaxChunkPerDoc int `json:"max_chunk_per_doc,omitempty"`
OverLapTokens int `json:"overlap_tokens,omitempty"`
}
func (r *RerankRequest) IsStream(c *gin.Context) bool {
return false
}
func (r *RerankRequest) GetTokenCountMeta() *types.TokenCountMeta {
var texts = make([]string, 0)
for _, document := range r.Documents {
texts = append(texts, fmt.Sprintf("%v", document))
}
if r.Query != "" {
texts = append(texts, r.Query)
}
return &types.TokenCountMeta{
CombineText: strings.Join(texts, "\n"),
}
}
func (r *RerankRequest) SetModelName(modelName string) {
if modelName != "" {
r.Model = modelName
}
}
func (r *RerankRequest) GetReturnDocuments() bool {
if r.ReturnDocuments == nil {
return false
}
return *r.ReturnDocuments
}
type RerankResponseResult struct {
Document any `json:"document,omitempty"`
Index int `json:"index"`
RelevanceScore float64 `json:"relevance_score"`
}
type RerankDocument struct {
Text any `json:"text"`
}
type RerankResponse struct {
Results []RerankResponseResult `json:"results"`
Usage Usage `json:"usage"`
}
|