File size: 1,323 Bytes
6b5e3dd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
import numpy as np
import torch
import torch.nn as nn
from transformers import Pipeline
class RegressionPipeline(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):
print(inputs)
encoded_corpus = self.tokenizer(text=inputs,
add_special_tokens=True,
padding='max_length',
truncation='longest_first',
max_length=300,
return_attention_mask=True)
return {"model_input": encoded_corpus}
def _forward(self, model_inputs):
print(model_inputs)
# model_inputs == {"model_input": model_input}
outputs = self.model(torch.tensor(model_inputs['model_input']['input_ids']).reshape(1, -1).to(torch.int64),
torch.tensor(model_inputs['model_input']['attention_mask']).reshape(1, -1).to(torch.int64))
return outputs
def postprocess(self, model_outputs):
print(model_outputs)
return model_outputs.numpy()
|