|
from transformers import Pipeline |
|
|
|
class MyPipeline(Pipeline): |
|
def _sanitize_parameters(self, **kwargs): |
|
preprocess_kwargs = {} |
|
if "max_length" in kwargs: |
|
preprocess_kwargs["max_length"] = kwargs["max_length"] |
|
if "num_beams" in kwargs: |
|
preprocess_kwargs["num_beams"] = kwargs["num_beams"] |
|
|
|
return preprocess_kwargs, {}, {} |
|
def preprocess(self, inputs, **kwargs): |
|
inputs = re.sub(r'[^A-Za-z가-힣,<>0-9:&# ]', '', inputs) |
|
inputs = "질문 생성: <unused0>"+inputs |
|
|
|
input_ids = [tokenizer.bos_token_id] + tokenizer.encode(inputs) + [tokenizer.eos_token_id] |
|
return {"inputs":torch.tensor([input_ids]),'max_length':kwargs['max_length'],'num_beams':kwargs['num_beams'] } |
|
|
|
def _forward(self, model_inputs): |
|
res_ids = model.generate( |
|
model_inputs['inputs'], |
|
max_length=model_inputs['max_length'], |
|
num_beams=model_inputs['num_beams'], |
|
eos_token_id=tokenizer.eos_token_id, |
|
bad_words_ids=[[tokenizer.unk_token_id]] |
|
) |
|
return {"logits": res_ids} |
|
|
|
def postprocess(self, model_outputs): |
|
a = tokenizer.batch_decode(model_outputs["logits"].tolist())[0] |
|
out_question = a.replace('<s>', '').replace('</s>', '') |
|
return out_question |
|
|
|
def _inference(self,paragraph,**kwargs): |
|
input_ids = self.preprocess(paragraph,**kwargs) |
|
reds_ids = self._forward(input_ids) |
|
out_question = self.postprocess(reds_ids) |
|
return out_question |
|
|
|
def make_question(self, text, **kwargs): |
|
words = text.split(" ") |
|
frame_size = kwargs['frame_size'] |
|
hop_length = kwargs['hop_length'] |
|
steps = round((len(words)-frame_size)/hop_length) + 1 |
|
outs = [] |
|
for step in range(steps): |
|
try: |
|
script = " ".join(words[step*hop_length:step*hop_length+frame_size]) |
|
except: |
|
script = " ".join(words[(1+step)*hop_length:]) |
|
|
|
outs.append(self._inference(script,**kwargs)) |
|
|
|
|
|
return outs |
