| | 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 |
