-
-
Notifications
You must be signed in to change notification settings - Fork 2.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(rerankers): Add new backend, support jina rerankers API
Signed-off-by: Ettore Di Giacinto <[email protected]>
- Loading branch information
Showing
19 changed files
with
358 additions
and
5 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| name: jina-reranker-v1-base-en | ||
| backend: rerankers | ||
| parameters: | ||
| model: cross-encoder | ||
|
|
||
| usage: | | ||
| You can test this model with curl like this: | ||
| curl http://localhost:8080/v1/rerank \ | ||
| -H "Content-Type: application/json" \ | ||
| -d '{ | ||
| "model": "jina-reranker-v1-base-en", | ||
| "query": "Organic skincare products for sensitive skin", | ||
| "documents": [ | ||
| "Eco-friendly kitchenware for modern homes", | ||
| "Biodegradable cleaning supplies for eco-conscious consumers", | ||
| "Organic cotton baby clothes for sensitive skin", | ||
| "Natural organic skincare range for sensitive skin", | ||
| "Tech gadgets for smart homes: 2024 edition", | ||
| "Sustainable gardening tools and compost solutions", | ||
| "Sensitive skin-friendly facial cleansers and toners", | ||
| "Organic food wraps and storage solutions", | ||
| "All-natural pet food for dogs with allergies", | ||
| "Yoga mats made from recycled materials" | ||
| ], | ||
| "top_n": 3 | ||
| }' |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| name: jina-reranker-v1-base-en | ||
| backend: rerankers | ||
| parameters: | ||
| model: cross-encoder | ||
|
|
||
| usage: | | ||
| You can test this model with curl like this: | ||
| curl http://localhost:8080/v1/rerank \ | ||
| -H "Content-Type: application/json" \ | ||
| -d '{ | ||
| "model": "jina-reranker-v1-base-en", | ||
| "query": "Organic skincare products for sensitive skin", | ||
| "documents": [ | ||
| "Eco-friendly kitchenware for modern homes", | ||
| "Biodegradable cleaning supplies for eco-conscious consumers", | ||
| "Organic cotton baby clothes for sensitive skin", | ||
| "Natural organic skincare range for sensitive skin", | ||
| "Tech gadgets for smart homes: 2024 edition", | ||
| "Sustainable gardening tools and compost solutions", | ||
| "Sensitive skin-friendly facial cleansers and toners", | ||
| "Organic food wraps and storage solutions", | ||
| "All-natural pet food for dogs with allergies", | ||
| "Yoga mats made from recycled materials" | ||
| ], | ||
| "top_n": 3 | ||
| }' |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,27 @@ | ||
| name: jina-reranker-v1-base-en | ||
| backend: rerankers | ||
| parameters: | ||
| model: cross-encoder | ||
|
|
||
| usage: | | ||
| You can test this model with curl like this: | ||
| curl http://localhost:8080/v1/rerank \ | ||
| -H "Content-Type: application/json" \ | ||
| -d '{ | ||
| "model": "jina-reranker-v1-base-en", | ||
| "query": "Organic skincare products for sensitive skin", | ||
| "documents": [ | ||
| "Eco-friendly kitchenware for modern homes", | ||
| "Biodegradable cleaning supplies for eco-conscious consumers", | ||
| "Organic cotton baby clothes for sensitive skin", | ||
| "Natural organic skincare range for sensitive skin", | ||
| "Tech gadgets for smart homes: 2024 edition", | ||
| "Sustainable gardening tools and compost solutions", | ||
| "Sensitive skin-friendly facial cleansers and toners", | ||
| "Organic food wraps and storage solutions", | ||
| "All-natural pet food for dogs with allergies", | ||
| "Yoga mats made from recycled materials" | ||
| ], | ||
| "top_n": 3 | ||
| }' |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,39 @@ | ||
| package backend | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
|
|
||
| "github.com/go-skynet/LocalAI/core/config" | ||
| "github.com/go-skynet/LocalAI/pkg/grpc/proto" | ||
| model "github.com/go-skynet/LocalAI/pkg/model" | ||
| ) | ||
|
|
||
| func Rerank(backend, modelFile string, request *proto.RerankRequest, loader *model.ModelLoader, appConfig *config.ApplicationConfig, backendConfig config.BackendConfig) (*proto.RerankResult, error) { | ||
| bb := backend | ||
| if bb == "" { | ||
| bb = model.PiperBackend | ||
| } | ||
|
|
||
| grpcOpts := gRPCModelOpts(backendConfig) | ||
|
|
||
| opts := modelOpts(config.BackendConfig{}, appConfig, []model.Option{ | ||
| model.WithBackendString(bb), | ||
| model.WithModel(modelFile), | ||
| model.WithContext(appConfig.Context), | ||
| model.WithAssetDir(appConfig.AssetsDestination), | ||
| model.WithLoadGRPCLoadModelOpts(grpcOpts), | ||
| }) | ||
| rerankModel, err := loader.BackendLoader(opts...) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| if rerankModel == nil { | ||
| return nil, fmt.Errorf("could not load piper model") | ||
| } | ||
|
|
||
| res, err := rerankModel.Rerank(context.Background(), request) | ||
|
|
||
| return res, err | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,84 @@ | ||
| package jina | ||
|
|
||
| import ( | ||
| "github.com/go-skynet/LocalAI/core/backend" | ||
| "github.com/go-skynet/LocalAI/core/config" | ||
|
|
||
| fiberContext "github.com/go-skynet/LocalAI/core/http/ctx" | ||
| "github.com/go-skynet/LocalAI/core/schema" | ||
| "github.com/go-skynet/LocalAI/pkg/grpc/proto" | ||
| "github.com/go-skynet/LocalAI/pkg/model" | ||
| "github.com/gofiber/fiber/v2" | ||
| "github.com/rs/zerolog/log" | ||
| ) | ||
|
|
||
| func JINARerankEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error { | ||
| return func(c *fiber.Ctx) error { | ||
| req := new(schema.JINARerankRequest) | ||
| if err := c.BodyParser(req); err != nil { | ||
| return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{ | ||
| "error": "Cannot parse JSON", | ||
| }) | ||
| } | ||
|
|
||
| input := new(schema.TTSRequest) | ||
|
|
||
| // Get input data from the request body | ||
| if err := c.BodyParser(input); err != nil { | ||
| return err | ||
| } | ||
|
|
||
| modelFile, err := fiberContext.ModelFromContext(c, ml, input.Model, false) | ||
| if err != nil { | ||
| modelFile = input.Model | ||
| log.Warn().Msgf("Model not found in context: %s", input.Model) | ||
| } | ||
|
|
||
| cfg, err := cl.LoadBackendConfigFileByName(modelFile, appConfig.ModelPath, | ||
| config.LoadOptionDebug(appConfig.Debug), | ||
| config.LoadOptionThreads(appConfig.Threads), | ||
| config.LoadOptionContextSize(appConfig.ContextSize), | ||
| config.LoadOptionF16(appConfig.F16), | ||
| ) | ||
|
|
||
| if err != nil { | ||
| modelFile = input.Model | ||
| log.Warn().Msgf("Model not found in context: %s", input.Model) | ||
| } else { | ||
| modelFile = cfg.Model | ||
| } | ||
| log.Debug().Msgf("Request for model: %s", modelFile) | ||
|
|
||
| if input.Backend != "" { | ||
| cfg.Backend = input.Backend | ||
| } | ||
|
|
||
| request := &proto.RerankRequest{ | ||
| Query: req.Query, | ||
| TopN: int32(req.TopN), | ||
| Documents: req.Documents, | ||
| } | ||
|
|
||
| results, err := backend.Rerank(cfg.Backend, modelFile, request, ml, appConfig, *cfg) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| response := &schema.JINARerankResponse{ | ||
| Model: req.Model, | ||
| } | ||
|
|
||
| for _, r := range results.Results { | ||
| response.Results = append(response.Results, schema.JINADocumentResult{ | ||
| Index: int(r.Index), | ||
| Document: schema.JINAText{Text: r.Text}, | ||
| RelevanceScore: float64(r.RelevanceScore), | ||
| }) | ||
| } | ||
|
|
||
| response.Usage.TotalTokens = int(results.Usage.TotalTokens) | ||
| response.Usage.PromptTokens = int(results.Usage.PromptTokens) | ||
|
|
||
| return c.Status(fiber.StatusOK).JSON(response) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,19 @@ | ||
| package routes | ||
|
|
||
| import ( | ||
| "github.com/go-skynet/LocalAI/core/config" | ||
| "github.com/go-skynet/LocalAI/core/http/endpoints/jina" | ||
|
|
||
| "github.com/go-skynet/LocalAI/pkg/model" | ||
| "github.com/gofiber/fiber/v2" | ||
| ) | ||
|
|
||
| func RegisterJINARoutes(app *fiber.App, | ||
| cl *config.BackendConfigLoader, | ||
| ml *model.ModelLoader, | ||
| appConfig *config.ApplicationConfig, | ||
| auth func(*fiber.Ctx) error) { | ||
|
|
||
| // POST endpoint to mimic the reranking | ||
| app.Post("/v1/rerank", jina.JINARerankEndpoint(cl, ml, appConfig)) | ||
| } |
Oops, something went wrong.