one-api/relay/adaptor/vertexai/embedding/model.go

46 lines
1.5 KiB
Go

package vertexai
type EmbeddingTaskType string
const (
EmbeddingTaskTypeRetrievalQuery EmbeddingTaskType = "RETRIEVAL_QUERY"
EmbeddingTaskTypeRetrievalDocument EmbeddingTaskType = "RETRIEVAL_DOCUMENT"
EmbeddingTaskTypeSemanticSimilarity EmbeddingTaskType = "SEMANTIC_SIMILARITY"
EmbeddingTaskTypeClassification EmbeddingTaskType = "CLASSIFICATION"
EmbeddingTaskTypeClustering EmbeddingTaskType = "CLUSTERING"
EmbeddingTaskTypeQuestionAnswering EmbeddingTaskType = "QUESTION_ANSWERING"
EmbeddingTaskTypeFactVerification EmbeddingTaskType = "FACT_VERIFICATION"
EmbeddingTaskTypeCodeRetrievalQuery EmbeddingTaskType = "CODE_RETRIEVAL_QUERY"
)
type EmbeddingRequest struct {
Instances []EmbeddingInstance `json:"instances"`
Parameters EmbeddingParams `json:"parameters"`
}
type EmbeddingInstance struct {
Content string `json:"content"`
TaskType EmbeddingTaskType `json:"task_type,omitempty"`
Title string `json:"title,omitempty"`
}
type EmbeddingParams struct {
AutoTruncate bool `json:"autoTruncate,omitempty"`
OutputDimensionality int `json:"outputDimensionality,omitempty"`
// Texts []string `json:"texts,omitempty"`
}
type EmbeddingResponse struct {
Predictions []struct {
Embeddings EmbeddingData `json:"embeddings"`
} `json:"predictions"`
}
type EmbeddingData struct {
Statistics struct {
Truncated bool `json:"truncated"`
TokenCount int `json:"token_count"`
} `json:"statistics"`
Values []float64 `json:"values"`
}