import os import json import tempfile import sys import datetime import re import string sys.path.append('mtool') import torch from model.model import Model from data.dataset import Dataset from config.params import Params from utility.initialize import initialize from data.batch import Batch from mtool.main import main as mtool_main from tqdm import tqdm class PredictionModel(torch.nn.Module): def __init__(self, checkpoint_path=os.path.join('models', 'checkpoint.bin'), default_mrp_path=os.path.join('models', 'default.mrp'), verbose=False): super().__init__() self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') self.checkpoint = torch.load('./models/checkpoint.bin', map_location=torch.device('cpu')) self.verbose = verbose self.args = Params().load_state_dict(self.checkpoint['params']) self.args.log_wandb = False self.args.training_data = default_mrp_path self.args.validation_data = default_mrp_path self.args.test_data = default_mrp_path self.args.only_train = False self.args.encoder = os.path.join('models', 'encoder') initialize(self.args, init_wandb=False) self.dataset = Dataset(self.args, verbose=False) self.model = Model(self.dataset, self.args).to(self.device) self.model.load_state_dict(self.checkpoint["model"], strict=False) self.model.eval() def _mrp_to_text(self, mrp_list, graph_mode='labeled-edge'): framework = 'norec' with tempfile.NamedTemporaryFile(delete=False, mode='w') as output_text_file: output_text_filename = output_text_file.name with tempfile.NamedTemporaryFile(delete=False, mode='w') as mrp_file: line = '\n'.join([json.dumps(entry) for entry in mrp_list]) mrp_file.write(line) mrp_filename = mrp_file.name if graph_mode == 'labeled-edge': mtool_main([ '--strings', '--ids', '--read', 'mrp', '--write', framework, mrp_filename, output_text_filename ]) elif graph_mode == 'node-centric': mtool_main([ '--node_centric', '--strings', '--ids', '--read', 'mrp', '--write', framework, mrp_filename, output_text_filename ]) else: raise Exception(f'Unknown graph mode: {graph_mode}') with open(output_text_filename) as f: texts = json.load(f) os.unlink(output_text_filename) os.unlink(mrp_filename) return texts def clean_texts(self, texts): punctuation = ''.join([f'\\{s}' for s in string.punctuation]) texts = [re.sub(f'([{punctuation}])', ' \\1 ', t) for t in texts] texts = [re.sub(r' +', ' ', t) for t in texts] return texts def _predict_to_mrp(self, texts, graph_mode='labeled-edge'): texts = self.clean_texts(texts) framework, language = self.args.framework, self.args.language data = self.dataset.load_sentences(texts, self.args) res_sentences = {f"{i}": {'input': sentence} for i, sentence in enumerate(texts)} date_str = datetime.datetime.now().date().isoformat() for key, value_dict in res_sentences.items(): value_dict['id'] = key value_dict['time'] = date_str value_dict['framework'], value_dict['language'] = framework, language value_dict['nodes'], value_dict['edges'], value_dict['tops'] = [], [], [] for i, batch in enumerate(tqdm(data) if self.verbose else data): with torch.no_grad(): predictions = self.model(Batch.to(batch, self.device), inference=True) for prediction in predictions: for key, value in prediction.items(): res_sentences[prediction['id']][key] = value return res_sentences def predict(self, text_list, graph_mode='labeled-edge', language='no'): mrp_predictions = self._predict_to_mrp(text_list, graph_mode) predictions = self._mrp_to_text(mrp_predictions.values(), graph_mode) return predictions def forward(self, x): return self.predict(x)