Feature/add jina reranker rebased

[emrgnt-cmplxty]

🚀 This description was created by Ellipsis for commit e8a612c

Summary:

This PR introduces a reranking feature to the R2R project, modifies various parameters related to search and reranking, updates configuration files, and modifies several classes to support the new reranking feature.

Key points:

• Added reranking feature to the R2R project.
• Modified search limit and rerank limit parameters in various functions and methods.
• Updated configuration files to support the new reranking feature.
• Modified SentenceTransformerEmbeddingProvider class to support separate models for search and reranking.
• Updated QnARAGPipeline and WebRAGPipeline classes to support the new reranking feature.

---

Generated with ❤️ by ellipsis.dev

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
r2r-docs	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Apr 19, 2024 11:41pm

[ellipsis-dev]

❌ Changes requested.

• Reviewed the entire pull request up to e8a612c
• Looked at 1684 lines of code in 26 files
• Took 7 minutes and 38 seconds to review

More info

• Skipped 1 files when reviewing.
• Skipped posting 3 additional comments because they didn't meet confidence threshold of 50%.

1. r2r/core/pipelines/rag.py:85:

• Assessed confidence : 100%
• Grade: 0%
• Comment:
  The abstract method rerank_results should also include the limit parameter to match the implementation in the subclasses.

    def rerank_results(
        self, transformed_query: Any, results: list, limit: int
    ) -> list:

• Reasoning:
  The rerank_results function in the RAGPipeline class has been updated to include a limit parameter, but the abstract method in the base class has not been updated to reflect this change. This could lead to confusion or errors when implementing new subclasses of RAGPipeline.

2. r2r/pipelines/web/rag.py:53:

• Assessed confidence : 80%
• Grade: 0%
• Comment:
  The run function should check if search_limit and rerank_limit are positive integers.

    def run(
        self,
        query,
        filters={},
        search_limit=25,
        rerank_limit=15,
        search_only=False,
        generation_config: Optional[GenerationConfig] = None,
        *args,
        **kwargs,
    ):
        if not isinstance(search_limit, int) or search_limit <= 0:
            raise ValueError("'search_limit' must be a positive integer")
        if not isinstance(rerank_limit, int) or rerank_limit <= 0:
            raise ValueError("'rerank_limit' must be a positive integer")
        self.initialize_pipeline(query, False)

        logger.debug(
            f"Running the `WebRAGPipeline` with pipeline_run_info={self.pipeline_run_info}."
        )

        transformed_query = self.transform_query(query)
        search_results = self.search(transformed_query, filters, limit)
        context = self.construct_context(search_results)
        prompt = self.construct_prompt({"query": query, "context": context})

        if search_only:
            return RAGPipelineOutput(search_results, None, None)
        elif not generation_config:
            return self.generate_completion(transformed_query, search_results)
        else:
            return self.generate_completion(
                transformed_query, search_results, generation_config
            )

• Reasoning:
  The run function in the WebRAGPipeline class has been updated to include search_limit and rerank_limit parameters, but the function does not check if these parameters are positive integers. This could lead to unexpected behavior or errors when the function is called with invalid search_limit or rerank_limit arguments.

3. r2r/pipelines/core/embedding.py:123:

• Assessed confidence : 80%
• Grade: 0%
• Comment:
  The run function should check if search_limit and rerank_limit are positive integers.

    def run(
        self,
        document: DocumentPage,
        do_chunking=False,
        do_upsert=True,
        **kwargs: Any,
    ):
        if not isinstance(search_limit, int) or search_limit <= 0:
            raise ValueError("'search_limit' must be a positive integer")
        if not isinstance(rerank_limit, int) or rerank_limit <= 0:
            raise ValueError("'rerank_limit' must be a positive integer")
        self.initialize_pipeline(document=document)

        logger.debug(
            f"Running the `BasicEmbeddingPipeline` with pipeline_run_info={self.pipeline_run_info}."
        )

        batch_data = []

        chunks = (
            self.chunk_text(document.text) if do_chunking else [document.text]
        )
        for chunk_iter, chunk in enumerate(chunks):
            batch_data.append(
                (
                    document.document_id,
                    document.page_number,
                    chunk_iter,
                    chunk,
                    copy.copy(document.metadata),
                )
            )

            if len(batch_data) == self.embedding_batch_size:
                self._process_batches(batch_data, do_upsert)
                batch_data = []

        # Process any remaining batch
        if batch_data:
            self._process_batches(batch_data, do_upsert)

• Reasoning:
  The run function in the BasicEmbeddingPipeline class has been updated to include search_limit and rerank_limit parameters, but the function does not check if these parameters are positive integers. This could lead to unexpected behavior or errors when the function is called with invalid search_limit or rerank_limit arguments.

Workflow ID: wflow_9lgHObXDBColvx9w

---

Want Ellipsis to fix these issues? Tag @ellipsis-dev in a comment. We'll respond in a few minutes. Learn more here.

[ellipsis-dev]

There is a typo in the parameter name. It should be rerank_limit instead of renark_limit.

Suggested change

	query.renark_limit,
	"rerank_limit": rerank_limit,

[ellipsis-dev]

The rerank_results function should check if limit is a positive integer.

Suggested change

	) -> list[VectorSearchResult]:
	def rerank_results(
	self, transformed_query: str, results: list[VectorSearchResult], limit: int
	) -> list[VectorSearchResult]:
	if not isinstance(limit, int) or limit <= 0:
	raise ValueError("'limit' must be a positive integer")
	return list(reversed(results))[0:limit]

[ellipsis-dev]

The search function should check if search_limit and rerank_limit are positive integers.

Suggested change

	def search(
	def search(
	self,
	query: str,
	search_limit: Optional[int] = 25,
	rerank_limit: Optional[int] = 15,
	filters: Optional[Dict[str, Any]] = None,
	settings: Optional[Dict[str, Any]] = None,
	):
	if not isinstance(search_limit, int) or search_limit <= 0:
	raise ValueError("'search_limit' must be a positive integer")
	if not isinstance(rerank_limit, int) or rerank_limit <= 0:
	raise ValueError("'rerank_limit' must be a positive integer")
	json_data = {
	"query": query,
	"filters": filters or {},
	"search_limit": search_limit,
	"rerank_limit": rerank_limit,
	"settings": settings or {},
	}
	response = requests.post(url, json=json_data)
	return response.json()

[ellipsis-dev]

The run function should check if search_limit and rerank_limit are positive integers.

Suggested change

	def run(
	self,
	query,
	filters={},
	search_limit=25,
	rerank_limit=15,
	search_only=False,
	generation_config: Optional[GenerationConfig] = None,
	*args,
	**kwargs,
	):
	if not isinstance(search_limit, int) or search_limit <= 0:
	raise ValueError("'search_limit' must be a positive integer")
	if not isinstance(rerank_limit, int) or rerank_limit <= 0:
	raise ValueError("'rerank_limit' must be a positive integer")
	self.initialize_pipeline(query, search_only)
	transformed_query = self.transform_query(query)
	search_results = self.search(transformed_query, filters, search_limit)
	search_results = self.rerank_results(
	transformed_query, search_results, rerank_limit
	)
	if search_only:
	return RAGPipelineOutput(search_results, None, None)
	elif not generation_config:
	return self.generate_completion(transformed_query, search_results)
	else:
	return self.generate_completion(
	transformed_query, search_results, generation_config
	)

[ellipsis-dev]

The get_embedding function should check if stage is an instance of PipelineStage.

Suggested change

	self.CrossEncoder = CrossEncoder
	def get_embedding(
	self, text: str, stage: PipelineStage = PipelineStage.SEARCH
	) -> list[float]:
	if not isinstance(stage, PipelineStage):
	raise ValueError("'stage' must be an instance of `PipelineStage`")
	if stage != PipelineStage.SEARCH:
	raise ValueError("`get_embedding` only supports `SEARCH` stage.")
	if not self.do_search:
	raise ValueError(
	"`get_embedding` can only be called for the search stage if a search model is set."
	)
	encoder = self.search_encoder
	return encoder.encode([text]).tolist()[0]