Method process_text_batch

src/refined/inference/processor.py:197–271

Batched version of process_text(). This method uses the convert_batch_elements_to_batched_tns() method which could be much more efficient by directly returning tensors that do not require indexing - like done in _identify_entity_mentions(). Currently, performance gai

(
            self,
            texts: List[str],
            spanss: Optional[List[List[Span]]] = None,
            ner_threshold: float = 0.5,
            prune_ner_types: bool = True,
            max_batch_size: int = 16,
            apply_class_check: bool = False,
            return_special_spans: bool = True,
            sort_by_tokens: bool = True
    )

        return all_spans

    def process_text_batch(
            self,
            texts: List[str],
            spanss: Optional[List[List[Span]]] = None,
            ner_threshold: float = 0.5,
            prune_ner_types: bool = True,
            max_batch_size: int = 16,
            apply_class_check: bool = False,
            return_special_spans: bool = True,
            sort_by_tokens: bool = True
    ) -> List[Doc]:
        """
        Batched version of process_text().
        This method uses the convert_batch_elements_to_batched_tns() method which could be much more
        efficient by directly returning tensors that do not require indexing - like done in _identify_entity_mentions().
        Currently, performance gain from batching is mostly present when document texts
        are short (such as question text).
        TODO: add optional already "indexed into" tensors to RefinedModel.forward() and use this during inference.
        # Note that "indexed into" tensors should not be used during multi-gpu training as it can cause issues.
        # `sort_by_tokens` is used to batch similar length chunks together for efficiency.
        # The current implementation means that the person name co-rereference trick will no longer work
        # across (512 token) chunks of the same document but there is a simple fix in _identify_entity_mentions().
        """
        all_spans = []

        docs = []
        if spanss is not None:
            for i, (text, spans) in enumerate(zip(texts, spanss)):
                doc = Doc.from_text_with_spans(
                    text, spans, self.preprocessor, backward_coref=self.backward_coref,
                    doc_id=i
                )
                docs.append(doc)
        else:
            for i, text in enumerate(texts):
                doc = Doc.from_text(
                    text, preprocessor=self.preprocessor, doc_id=i
                )
                docs.append(doc)
        batch_elements = [elem for doc in docs for elem in doc.to_batch_elements(
            preprocessor=self.preprocessor, override_max_seq=self.max_seq
        )]
        tns: Iterable[BatchedElementsTns] = convert_batch_elements_to_batched_tns(
                batch_elements,
                self.preprocessor,
                max_batch_size=max_batch_size,
                sort_by_tokens=sort_by_tokens,
            # TODO: currently sort_by_tokens=True means some person name co-reference can be missed.
            # To make this equivalent to running with sort_by_tokens=False, adjust the
            # in _identify_entity_mentions() is adjusted to keep track of names per doc and not assume batch_elements
            # for each document are sequential.
        )

        self.model.eval()

        for batch_idx, batch in enumerate(tns):
            batch_spans = self.process_tensors(batch=batch, ner_threshold=ner_threshold,
                                               return_special_spans=return_special_spans)