feat(rerankers): Add new backend, support jina rerankers API (#2121)

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

core/backend/rerank.go

@@ -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 == "" {
+		return nil, fmt.Errorf("backend is required")
+	}
+
+	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 rerank model")
+	}
+
+	res, err := rerankModel.Rerank(context.Background(), request)
+
+	return res, err
+}

core/http/app.go

@@ -194,6 +194,7 @@ func App(cl *config.BackendConfigLoader, ml *model.ModelLoader, appConfig *confi
 	routes.RegisterOpenAIRoutes(app, cl, ml, appConfig, auth)
 	routes.RegisterPagesRoutes(app, cl, ml, appConfig, auth)
 	routes.RegisterUIRoutes(app, cl, ml, appConfig, galleryService, auth)
+	routes.RegisterJINARoutes(app, cl, ml, appConfig, auth)
 
 	// Define a custom 404 handler
 	// Note: keep this at the bottom!

core/http/endpoints/jina/rerank.go

@@ -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)
+	}
+}

core/http/routes/jina.go

@@ -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))
+}

core/schema/jina.go

@@ -0,0 +1,34 @@
+package schema
+
+// RerankRequest defines the structure of the request payload
+type JINARerankRequest struct {
+	Model     string   `json:"model"`
+	Query     string   `json:"query"`
+	Documents []string `json:"documents"`
+	TopN      int      `json:"top_n"`
+}
+
+// DocumentResult represents a single document result
+type JINADocumentResult struct {
+	Index          int      `json:"index"`
+	Document       JINAText `json:"document"`
+	RelevanceScore float64  `json:"relevance_score"`
+}
+
+// Text holds the text of the document
+type JINAText struct {
+	Text string `json:"text"`
+}
+
+// RerankResponse defines the structure of the response payload
+type JINARerankResponse struct {
+	Model   string               `json:"model"`
+	Usage   JINAUsageInfo        `json:"usage"`
+	Results []JINADocumentResult `json:"results"`
+}
+
+// UsageInfo holds information about usage of tokens
+type JINAUsageInfo struct {
+	TotalTokens  int `json:"total_tokens"`
+	PromptTokens int `json:"prompt_tokens"`
+}