def postprocess()

in optimum/neuron/pipelines/transformers/sentence_transformers.py [0:0]

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

    def postprocess(self, _model_outputs, return_tensors=False):
        # Needed change for sentence transformers.
        # Check if the model outputs sentence embeddings or not.
        if hasattr(_model_outputs, "sentence_embedding"):
            model_outputs = _model_outputs.sentence_embedding
        else:
            model_outputs = _model_outputs
        # [0] is the first available tensor, logits or last_hidden_state.
        if return_tensors:
            return model_outputs[0]
        if self.framework == "pt":
            return model_outputs[0].tolist()