| thon | |
| from transformers import Pipeline | |
| class MyPipeline(Pipeline): | |
| def _sanitize_parameters(self, **kwargs): | |
| preprocess_kwargs = {} | |
| if "maybe_arg" in kwargs: | |
| preprocess_kwargs["maybe_arg"] = kwargs["maybe_arg"] | |
| return preprocess_kwargs, {}, {} | |
| def preprocess(self, inputs, maybe_arg=2): | |
| model_input = Tensor(inputs["input_ids"]) | |
| return {"model_input": model_input} | |
| def _forward(self, model_inputs): | |
| # model_inputs == {"model_input": model_input} | |
| outputs = self.model(**model_inputs) | |
| # Maybe {"logits": Tensor()} | |
| return outputs | |
| def postprocess(self, model_outputs): | |
| best_class = model_outputs["logits"].softmax(-1) | |
| return best_class | |
| The structure of this breakdown is to support relatively seamless support for CPU/GPU, while supporting doing | |
| pre/postprocessing on the CPU on different threads | |
| preprocess will take the originally defined inputs, and turn them into something feedable to the model. |