| from transformers import Pipeline, AutoTokenizer | |
| from torch_geometric.data import Batch | |
| import torch | |
| class GRetrieverPipeline(Pipeline): | |
| def __init__(self, **kwargs): | |
| Pipeline.__init__(self, **kwargs) | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model.config._name_or_path) | |
| self.eos_user = "<|eot_id|><|start_header_id|>assistant<|end_header_id|>" | |
| self.max_txt_len = self.model.config.max_txt_len | |
| self.bos_length = len(self.model.config.bos_id) | |
| self.input_length = 0 | |
| def _sanitize_parameters(self, **kwargs): | |
| preprocess_kwargs = {} | |
| if "textualized_graph" in kwargs: | |
| preprocess_kwargs["textualized_graph"] = kwargs["textualized_graph"] | |
| if "graph" in kwargs: | |
| preprocess_kwargs["graph"] = kwargs["graph"] | |
| if "generate_kwargs" in kwargs: | |
| preprocess_kwargs["generate_kwargs"] = kwargs["generate_kwargs"] | |
| return preprocess_kwargs, {}, {} | |
| def preprocess(self, inputs, textualized_graph, graph, generate_kwargs=None): | |
| textualized_graph_ids = self.tokenizer(textualized_graph, add_special_tokens=False)["input_ids"][:self.max_txt_len] | |
| question_ids = self.tokenizer(inputs, add_special_tokens=False)["input_ids"] | |
| eos_user_ids = self.tokenizer(self.eos_user, add_special_tokens=False)["input_ids"] | |
| input_ids = torch.tensor([ | |
| [-1]*(self.bos_length + 1) | |
| + textualized_graph_ids | |
| + question_ids | |
| + eos_user_ids | |
| ]) | |
| model_inputs = { | |
| "input_ids": input_ids, | |
| "attention_mask": torch.ones_like(input_ids) | |
| } | |
| model_inputs.update({ | |
| "graph": Batch.from_data_list([graph]) | |
| }) | |
| if generate_kwargs != None: | |
| model_inputs.update(generate_kwargs) | |
| self.input_length = input_ids.shape[1] | |
| return model_inputs | |
| def _forward(self, model_inputs): | |
| model_outputs = self.model.generate(**model_inputs) | |
| return model_outputs | |
| def postprocess(self, model_outputs): | |
| return self.tokenizer.decode(model_outputs[0, self.input_length:]) |