def predict()

in backends/python/server/text_embeddings_server/models/classification_model.py [0:0]

• 9 lines of code
• 4 McCabe index (conditional complexity)

    def predict(self, batch: PaddedBatch) -> List[Score]:
        kwargs = {"input_ids": batch.input_ids, "attention_mask": batch.attention_mask}
        if self.has_token_type_ids:
            kwargs["token_type_ids"] = batch.token_type_ids
        if self.has_position_ids:
            kwargs["position_ids"] = batch.position_ids

        output = self.model(**kwargs, return_dict=True)
        all_scores = output.logits.tolist()
        return [Score(values=scores) for scores in all_scores]