Spaces:
Build error
Build error
File size: 7,998 Bytes
bbd6ee6 850c669 bbd6ee6 850c669 bbd6ee6 9bf6109 bbd6ee6 | 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | import os
import re
import logging
import pandas as pd
import torch
import transformers
from transformers import AutoTokenizer, AutoConfig, T5ForConditionalGeneration, BartForConditionalGeneration
TASK_PREFIX = {
"ve": "extract values",
"ag": "generate attributes",
"avg": "generate attributes and values",
"av": "attribute value extraction"
}
ADDITIONAL_SP_TOKENS = {'hl': '<hl>'}
def load_language_model(model_id,
):
""" Load language model from hugging face hub. """
tokenizer = AutoTokenizer.from_pretrained(model_id)
config = AutoConfig.from_pretrained(model_id)
# model class
if config.model_type == 't5':
model_class = T5ForConditionalGeneration.from_pretrained
elif config.model_type == 'bart':
model_class = BartForConditionalGeneration.from_pretrained
else:
raise ValueError(f"Unsupported model type: {config.model_type}")
param = {'config': config}
model = model_class(model_id, **param)
return tokenizer, model, config
class TransformersAVG:
""" Transformers Language Model for Attribute Value Generation. """
def __init__(self,
model: str = None,
max_input_length: int = 512,
max_target_length: int = 256,
model_ve: str = None,
max_target_length_ve: int = 34,
is_ag: bool = None,
is_avg: bool = None,
is_ve: bool = None,
is_av: bool = None
) -> None:
self.is_ag = 'ag' in model.split('-') if is_ag is None else is_ag
self.is_ve = 've' in model.split('-') if is_ve is None else is_ve
self.is_av = 'mlt' in model.split('-') if is_av is None else is_av
self.is_avg = 'end2end' in model.split('-') if is_avg is None else is_avg
self.model_name = model
self.max_input_length = max_input_length
self.max_target_length = max_target_length
self.model_name_ve = model_ve
self.max_target_length_ve = max_target_length_ve
# load model
self.tokenizer, self.model, config = load_language_model(
self.model_name,
)
# Setup GPU device
self.device = 'cuda' if torch.cuda.device_count() > 0 else 'cpu'
self.model.to(self.device)
def generate_av_end2end(self,
context: str,
num_beams: int = 4,
splitting_symbol: str = '|'
):
""" Generate attribute value pairs in an end2end fashion. """
logging.info(f"running model for 'attribute value pair generation'.")
model_input = self.tokenizer(context, max_length=self.max_input_length, truncation=True,
padding="max_length")
model_input = {k:torch.unsqueeze(torch.tensor(v),dim=0) for k,v in model_input.items()}
outputs = self.model.generate(**model_input, num_beams=num_beams, do_sample=True, max_length=self.max_target_length)
outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
predictions = self.__format_av(outputs)
return predictions
def generate_av_pipeline(self,
context: str,
num_beams: int = 4
):
""" Generate attribute value pairs using the pipeline: first extract values then generate attributes """
logging.info(f"running model for value candidate extraction.")
# load ve model
self.ve_tokenizer, self.ve_model, ve_config = load_language_model(
self.model_name_ve,
)
# generate values
model_input = self.ve_tokenizer(context, max_length=self.max_input_length, truncation=True,
padding="max_length")
model_input = {k:torch.unsqueeze(torch.tensor(v),dim=0) for k,v in model_input.items()}
outputs = self.ve_model.generate(**model_input, num_beams=num_beams, do_sample=True, max_length=self.max_target_length_ve)
outputs = self.ve_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
list_of_values = self.__format_v(outputs)
# generate attributes
list_of_attributes = []
logging.info(f"running model for attribute generation.")
list_of_contexts = self.__highlight_value(context, list_of_values)
model_input = self.tokenizer(list_of_contexts, max_length=self.max_input_length, truncation=True,
padding="max_length")
model_input = {k:torch.tensor(v) for k,v in model_input.items()}
outputs = self.model.generate(**model_input, num_beams=num_beams, do_sample=True, max_length=self.max_target_length)
outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
list_of_attributes.append(outputs)
list_of_attributes = [pred for preds in list_of_attributes for pred in preds]
return [(att, val) for att, val in zip(list_of_attributes, list_of_values)]
def generate_av_mul(self,
context: str,
num_beams: int = 4):
""" Generate attribute-value pairs using a multi-task approach: same model for attribute genration and value extraction """
logging.info(f"running model for value extraction.")
model_input = self.tokenizer(f"extract value {context}", max_length=self.max_input_length, truncation=True,
padding="max_length")
model_input = {k:torch.unsqueeze(torch.tensor(v),dim=0) for k,v in model_input.items()}
outputs = self.model.generate(**model_input, num_beams=num_beams, do_sample=True, max_length=self.max_target_length)
outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
list_of_values = self.__format_v(outputs)
# generate attributes
list_of_attributes = []
logging.info(f"running model for attribute generation.")
list_of_contexts = self.__highlight_value(context, list_of_values, prefix=True)
model_input = self.tokenizer(list_of_contexts, max_length=self.max_input_length, truncation=True,
padding="max_length")
model_input = {k:torch.tensor(v) for k,v in model_input.items()}
outputs = self.model.generate(**model_input, num_beams=num_beams, do_sample=True, max_length=self.max_target_length)
outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
list_of_attributes.append(outputs)
list_of_attributes = [pred for preds in list_of_attributes for pred in preds]
return [(att, val) for att, val in zip(list_of_attributes, list_of_values)]
def __highlight_value(self,
context: str,
list_of_values: list,
prefix: bool = False
):
list_of_contexts = []
for value in list_of_values:
my_context = context.replace(value, f'<hl> {value} <hl>')
if prefix:
my_context = f"generate attribute {my_context}"
list_of_contexts.append(my_context)
return list_of_contexts
def __format_v(self, predictions: str, splitting_symbol: str = '|') -> list:
my_list = []
for raw_string in predictions.split(splitting_symbol):
v = raw_string.strip()
my_list.append(v)
return my_list
def __format_av(self, predictions: str, splitting_symbol: str = '|') -> list:
my_list = []
for raw_string in predictions.split(splitting_symbol):
a = raw_string.replace('attribute: ', '')
a = a.split(',')[0].strip()
v = raw_string.replace('value: ', '')
v = v.split(',')[1].strip()
my_list.append((a,v))
return my_list |