def __init__(
self, model_type, model_name, args=None, use_cuda=True, cuda_device=-1, **kwargs,
):
"""
Initializes a RepresentationModel model.
Args:
model_type: The type of model (bert, roberta, gpt2)
model_name: The exact architecture and trained weights to use. This may be a Hugging Face Transformers compatible pre-trained model, a community model, or the path to a directory containing model files.
args (optional): Default args will be used if this parameter is not provided. If provided, it should be a dict containing the args that should be changed in the default args.
use_cuda (optional): Use GPU if available. Setting to False will force model to use CPU only.
cuda_device (optional): Specific GPU that should be used. Will use the first available GPU by default.
**kwargs (optional): For providing proxies, force_download, resume_download, cache_dir and other options specific to the 'from_pretrained' implementation where this will be supplied.
""" # noqa: ignore flake8"
MODEL_CLASSES = {
"bert": (BertConfig, BertForTextRepresentation, BertTokenizer),
"roberta": (RobertaConfig, BertForTextRepresentation, RobertaTokenizer),
"gpt2": (GPT2Config, GPT2ForTextRepresentation, GPT2Tokenizer),
}
self.args = self._load_model_args(model_name)
if isinstance(args, dict):
self.args.update_from_dict(args)
elif isinstance(args, ModelArgs):
self.args = args
if "sweep_config" in kwargs:
self.is_sweeping = True
sweep_config = kwargs.pop("sweep_config")
sweep_values = sweep_config_to_sweep_values(sweep_config)
self.args.update_from_dict(sweep_values)
else:
self.is_sweeping = False
if self.args.manual_seed:
random.seed(self.args.manual_seed)
np.random.seed(self.args.manual_seed)
torch.manual_seed(self.args.manual_seed)
if self.args.n_gpu > 0:
torch.cuda.manual_seed_all(self.args.manual_seed)
config_class, model_class, tokenizer_class = MODEL_CLASSES[model_type]
self.config = config_class.from_pretrained(model_name, **self.args.config)
if use_cuda:
if torch.cuda.is_available():
if cuda_device == -1:
self.device = torch.device("cuda")
else:
self.device = torch.device(f"cuda:{cuda_device}")
else:
raise ValueError(
"'use_cuda' set to True when cuda is unavailable."
" Make sure CUDA is available or set use_cuda=False."
)
else:
self.device = "cpu"
self.model = model_class.from_pretrained(model_name, config=self.config, **kwargs)
self.results = {}
if not use_cuda:
self.args.fp16 = False
self.tokenizer = tokenizer_class.from_pretrained(model_name, do_lower_case=self.args.do_lower_case, **kwargs)
self.args.model_name = model_name
self.args.model_type = model_type
if self.args.wandb_project and not wandb_available:
warnings.warn("wandb_project specified but wandb is not available. Wandb disabled.")
self.args.wandb_project = None
if self.args.model_type == "gpt2":
# should we add a custom tokenizer for this model?
self.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
self.model.resize_token_embeddings(len(self.tokenizer))
def __init__(
self,
model_type,
model_name,
args=None,
use_cuda=True,
cuda_device=-1,
**kwargs,
):
"""
Initializes a LanguageGenerationModel model.
Args:
model_type: The type of model (gpt2, ctrl, openai-gpt, xlnet, transfo-xl, xlm)
model_name: Default Transformer model name or path to a directory containing Transformer model file (pytorch_nodel.bin).
args (optional): Default args will be used if this parameter is not provided. If provided, it should be a dict containing the args that should be changed in the default args.
use_cuda (optional): Use GPU if available. Setting to False will force model to use CPU only.
cuda_device (optional): Specific GPU that should be used. Will use the first available GPU by default.
**kwargs (optional): For providing proxies, force_download, resume_download, cache_dir and other options specific to the 'from_pretrained' implementation where this will be supplied.
""" # noqa: ignore flake8"
MODEL_CLASSES = {
"gpt2": (GPT2Config, GPT2LMHeadModel, GPT2Tokenizer),
"ctrl": (CTRLConfig, CTRLLMHeadModel, CTRLTokenizer),
"openai-gpt": (OpenAIGPTConfig, OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
"xlnet": (XLNetConfig, XLNetLMHeadModel, XLNetTokenizer),
"transfo-xl": (TransfoXLConfig, TransfoXLLMHeadModel, TransfoXLTokenizer),
"xlm": (XLMConfig, XLMWithLMHeadModel, XLMTokenizer),
}
self.args = self._load_model_args(model_name)
if isinstance(args, dict):
self.args.update_from_dict(args)
elif isinstance(args, LanguageGenerationArgs):
self.args = args
if "sweep_config" in kwargs:
self.is_sweeping = True
sweep_config = kwargs.pop("sweep_config")
sweep_values = sweep_config_to_sweep_values(sweep_config)
self.args.update_from_dict(sweep_values)
else:
self.is_sweeping = False
if self.args.manual_seed:
random.seed(self.args.manual_seed)
np.random.seed(self.args.manual_seed)
torch.manual_seed(self.args.manual_seed)
if self.args.n_gpu > 0:
torch.cuda.manual_seed_all(self.args.manual_seed)
if use_cuda:
if torch.cuda.is_available():
if cuda_device == -1:
self.device = torch.device("cuda")
else:
self.device = torch.device(f"cuda:{cuda_device}")
else:
raise ValueError(
"'use_cuda' set to True when cuda is unavailable."
"Make sure CUDA is available or set `use_cuda=False`."
)
else:
self.device = "cpu"
if self.args.special_tokens_list:
self.tokenizer.add_tokens(
self.args.special_tokens_list, special_tokens=True
)
self.model.resize_token_embeddings(len(self.tokenizer))
self.args.model_name = model_name
self.args.model_type = model_type
config_class, model_class, tokenizer_class = MODEL_CLASSES[model_type]
if self.args.tokenizer_name:
self.tokenizer = tokenizer_class.from_pretrained(
self.args.tokenizer_name, cache_dir=self.args.cache_dir
)
else:
self.tokenizer = tokenizer_class.from_pretrained(
model_name, cache_dir=self.args.cache_dir, **kwargs
)
self.args.tokenizer_name = model_name
if self.args.config_name:
self.config = config_class.from_pretrained(
self.args.config_name, cache_dir=self.args.cache_dir
)
else:
self.config = config_class.from_pretrained(
model_name, cache_dir=self.args.cache_dir, **kwargs
)
self.model = model_class.from_pretrained(
model_name,
config=self.config,
cache_dir=self.args.cache_dir,
**kwargs,
)
self.model.to(self.device)
def __init__(
self,
model_type,
model_name,
num_labels=None,
pos_weight=None,
args=None,
use_cuda=True,
cuda_device=-1,
**kwargs,
):
"""
Initializes a MultiLabelClassification model.
Args:
model_type: The type of model (bert, roberta)
model_name: Default Transformer model name or path to a directory containing Transformer model file (pytorch_nodel.bin).
num_labels (optional): The number of labels or classes in the dataset.
pos_weight (optional): A list of length num_labels containing the weights to assign to each label for loss calculation.
args (optional): Default args will be used if this parameter is not provided. If provided, it should be a dict containing the args that should be changed in the default args.
use_cuda (optional): Use GPU if available. Setting to False will force model to use CPU only.
cuda_device (optional): Specific GPU that should be used. Will use the first available GPU by default.
**kwargs (optional): For providing proxies, force_download, resume_download, cache_dir and other options specific to the 'from_pretrained' implementation where this will be supplied.
""" # noqa: ignore flake8"
MODEL_CLASSES = {
"albert": (
AlbertConfig,
AlbertForMultiLabelSequenceClassification,
AlbertTokenizer,
),
"bert": (
BertConfig,
BertForMultiLabelSequenceClassification,
BertTokenizer,
),
"camembert": (
CamembertConfig,
CamembertForMultiLabelSequenceClassification,
CamembertTokenizer,
),
"distilbert": (
DistilBertConfig,
DistilBertForMultiLabelSequenceClassification,
DistilBertTokenizer,
),
"electra":
(ElectraConfig, ElectraForMultiLabelSequenceClassification,
ElectraTokenizer),
"flaubert": (
FlaubertConfig,
FlaubertForMultiLabelSequenceClassification,
FlaubertTokenizer,
),
"longformer": (
LongformerConfig,
LongformerForMultiLabelSequenceClassification,
LongformerTokenizer,
),
"roberta": (
RobertaConfig,
RobertaForMultiLabelSequenceClassification,
RobertaTokenizer,
),
"xlm":
(XLMConfig, XLMForMultiLabelSequenceClassification, XLMTokenizer),
"xlmroberta": (
XLMRobertaConfig,
XLMRobertaForMultiLabelSequenceClassification,
XLMRobertaTokenizer,
),
"xlnet": (
XLNetConfig,
XLNetForMultiLabelSequenceClassification,
XLNetTokenizer,
),
}
self.args = self._load_model_args(model_name)
if isinstance(args, dict):
self.args.update_from_dict(args)
elif isinstance(args, MultiLabelClassificationArgs):
self.args = args
if self.args.thread_count:
torch.set_num_threads(self.args.thread_count)
if "sweep_config" in kwargs:
self.is_sweeping = True
sweep_config = kwargs.pop("sweep_config")
sweep_values = sweep_config_to_sweep_values(sweep_config)
self.args.update_from_dict(sweep_values)
else:
self.is_sweeping = False
if self.args.manual_seed:
random.seed(self.args.manual_seed)
np.random.seed(self.args.manual_seed)
torch.manual_seed(self.args.manual_seed)
if self.args.n_gpu > 0:
torch.cuda.manual_seed_all(self.args.manual_seed)
if not use_cuda:
self.args.fp16 = False
config_class, model_class, tokenizer_class = MODEL_CLASSES[model_type]
if num_labels:
self.config = config_class.from_pretrained(model_name,
num_labels=num_labels,
**self.args.config)
self.num_labels = num_labels
else:
self.config = config_class.from_pretrained(model_name,
**self.args.config)
self.num_labels = self.config.num_labels
self.pos_weight = pos_weight
if use_cuda:
if torch.cuda.is_available():
if cuda_device == -1:
self.device = torch.device("cuda")
else:
self.device = torch.device(f"cuda:{cuda_device}")
else:
raise ValueError(
"'use_cuda' set to True when cuda is unavailable."
" Make sure CUDA is available or set use_cuda=False.")
else:
self.device = "cpu"
if not self.args.quantized_model:
if self.pos_weight:
self.model = model_class.from_pretrained(
model_name,
config=self.config,
pos_weight=torch.Tensor(self.pos_weight).to(self.device),
**kwargs)
else:
self.model = model_class.from_pretrained(model_name,
config=self.config,
**kwargs)
else:
quantized_weights = torch.load(
os.path.join(model_name, "pytorch_model.bin"))
if self.pos_weight:
self.model = model_class.from_pretrained(
None,
config=self.config,
state_dict=quantized_weights,
weight=torch.Tensor(self.pos_weight).to(self.device),
)
else:
self.model = model_class.from_pretrained(
None, config=self.config, state_dict=quantized_weights)
if self.args.dynamic_quantize:
self.model = torch.quantization.quantize_dynamic(self.model,
{torch.nn.Linear},
dtype=torch.qint8)
if self.args.quantized_model:
self.model.load_state_dict(quantized_weights)
if self.args.dynamic_quantize:
self.args.quantized_model = True
self.results = {}
self.tokenizer = tokenizer_class.from_pretrained(
model_name, do_lower_case=self.args.do_lower_case, **kwargs)
self.args.model_name = model_name
self.args.model_type = model_type
if self.args.wandb_project and not wandb_available:
warnings.warn(
"wandb_project specified but wandb is not available. Wandb disabled."
)
self.args.wandb_project = None
self.weight = None # Not implemented for multilabel
def __init__(
self, model_type, model_name, args=None, use_cuda=True, cuda_device=-1, **kwargs,
):
"""
Initializes a ClassificationModel model.
Args:
model_type: The type of model (gpt, gpt2)
model_name: Default Transformer model name or path to a directory containing Transformer model file (pytorch_nodel.bin).
args (optional): Default args will be used if this parameter is not provided. If provided, it should be a dict containing the args that should be changed in the default args.
use_cuda (optional): Use GPU if available. Setting to False will force model to use CPU only.
cuda_device (optional): Specific GPU that should be used. Will use the first available GPU by default.
**kwargs (optional): For providing proxies, force_download, resume_download, cache_dir and other options specific to the 'from_pretrained' implementation where this will be supplied.
""" # noqa: ignore flake8"
MODEL_CLASSES = {
"gpt": (OpenAIGPTConfig, OpenAIGPTDoubleHeadsModel, OpenAIGPTTokenizer),
"gpt2": (GPT2Config, GPT2DoubleHeadsModel, GPT2Tokenizer),
}
self.args = self._load_model_args(model_name)
if isinstance(args, dict):
self.args.update_from_dict(args)
elif isinstance(args, ConvAIArgs):
self.args = args
if "sweep_config" in kwargs:
sweep_config = kwargs.pop("sweep_config")
sweep_values = sweep_config_to_sweep_values(sweep_config)
self.args.update_from_dict(sweep_values)
if self.args.manual_seed:
random.seed(self.args.manual_seed)
np.random.seed(self.args.manual_seed)
torch.manual_seed(self.args.manual_seed)
if self.args.n_gpu > 0:
torch.cuda.manual_seed_all(self.args.manual_seed)
if not use_cuda:
self.args.fp16 = False
config_class, model_class, tokenizer_class = MODEL_CLASSES[model_type]
self.__dict__.update(kwargs)
self.config = config_class.from_pretrained(model_name, **self.args.config)
if not self.args.quantized_model:
self.model = model_class.from_pretrained(model_name, **kwargs)
else:
quantized_weights = torch.load(os.path.join(model_name, "pytorch_model.bin"))
self.model = model_class.from_pretrained(None, config=self.config, state_dict=quantized_weights)
self.tokenizer = tokenizer_class.from_pretrained(model_name, **kwargs)
self.add_special_tokens_(self.model, self.tokenizer)
if self.args.dynamic_quantize:
self.model = torch.quantization.quantize_dynamic(self.model, {torch.nn.Linear}, dtype=torch.qint8)
if self.args.quantized_model:
self.model.load_state_dict(quantized_weights)
if self.args.dynamic_quantize:
self.args.quantized_model = True
self.results = {}
if use_cuda:
if torch.cuda.is_available():
if cuda_device == -1:
self.device = torch.device("cuda")
else:
self.device = torch.device(f"cuda:{cuda_device}")
else:
raise ValueError(
"'use_cuda' set to True when cuda is unavailable."
" Make sure CUDA is available or set use_cuda=False."
)
else:
self.device = "cpu"
self.args.model_name = model_name
self.args.model_type = model_type
if self.args.wandb_project and not wandb_available:
warnings.warn("wandb_project specified but wandb is not available. Wandb disabled.")
self.args.wandb_project = None
def __init__(
self,
model_name,
args=None,
use_cuda=True,
cuda_device=-1,
**kwargs,
):
"""
Initializes a T5Model model.
Args:
model_name: The exact architecture and trained weights to use. This may be a Hugging Face Transformers compatible pre-trained model, a community model, or the path to a directory containing model files.
args (optional): Default args will be used if this parameter is not provided. If provided, it should be a dict containing the args that should be changed in the default args.
use_cuda (optional): Use GPU if available. Setting to False will force model to use CPU only.
cuda_device (optional): Specific GPU that should be used. Will use the first available GPU by default.
**kwargs (optional): For providing proxies, force_download, resume_download, cache_dir and other options specific to the 'from_pretrained' implementation where this will be supplied.
""" # noqa: ignore flake8"
self.args = self._load_model_args(model_name)
if isinstance(args, dict):
self.args.update_from_dict(args)
elif isinstance(args, T5Args):
self.args = args
if "sweep_config" in kwargs:
sweep_config = kwargs.pop("sweep_config")
sweep_values = sweep_config_to_sweep_values(sweep_config)
self.args.update_from_dict(sweep_values)
if self.args.manual_seed:
random.seed(self.args.manual_seed)
np.random.seed(self.args.manual_seed)
torch.manual_seed(self.args.manual_seed)
if self.args.n_gpu > 0:
torch.cuda.manual_seed_all(self.args.manual_seed)
if use_cuda:
if torch.cuda.is_available():
if cuda_device == -1:
self.device = torch.device("cuda")
else:
self.device = torch.device(f"cuda:{cuda_device}")
else:
raise ValueError(
"'use_cuda' set to True when cuda is unavailable."
"Make sure CUDA is available or set `use_cuda=False`.")
else:
self.device = "cpu"
self.results = {}
self.config = T5Config.from_pretrained(model_name, **self.args.config)
self.model = T5ForConditionalGeneration.from_pretrained(
model_name, config=self.config)
self.tokenizer = T5Tokenizer.from_pretrained(model_name, truncate=True)
if self.args.dynamic_quantize:
self.model = torch.quantization.quantize_dynamic(self.model,
{torch.nn.Linear},
dtype=torch.qint8)
if not use_cuda:
self.args.fp16 = False
self.args.model_type = "T5"
self.args.model_name = model_name
if self.args.wandb_project and not wandb_available:
warnings.warn(
"wandb_project specified but wandb is not available. Wandb disabled."
)
self.args.wandb_project = None
def __init__(
self,
encoder_type=None,
encoder_name=None,
decoder_name=None,
encoder_decoder_type=None,
encoder_decoder_name=None,
additional_special_tokens_encoder=None,
additional_special_tokens_decoder=None,
config=None,
args=None,
use_cuda=True,
cuda_device=-1,
**kwargs,
):
"""
Initializes a Seq2SeqModel.
Args:
encoder_type (optional): The type of model to use as the encoder.
encoder_name (optional): The exact architecture and trained weights to use. This may be a Hugging Face Transformers compatible pre-trained model, a community model, or the path to a directory containing model files.
decoder_name (optional): The exact architecture and trained weights to use. This may be a Hugging Face Transformers compatible pre-trained model, a community model, or the path to a directory containing model files.
Must be the same "size" as the encoder model (base/base, large/large, etc.)
encoder_decoder_type (optional): The type of encoder-decoder model. (E.g. bart)
encoder_decoder_name (optional): The path to a directory containing the saved encoder and decoder of a Seq2SeqModel. (E.g. "outputs/") OR a valid BART or MarianMT model.
additional_special_tokens_encoder (optional): dict of special tokens to add to encoder tokenizer
additional_special_tokens_decoder (optional): dict of special tokens to add to decoder tokenizer
config (optional): A configuration file to build an EncoderDecoderModel.
args (optional): Default args will be used if this parameter is not provided. If provided, it should be a dict containing the args that should be changed in the default args.
use_cuda (optional): Use GPU if available. Setting to False will force model to use CPU only.
cuda_device (optional): Specific GPU that should be used. Will use the first available GPU by default.
**kwargs (optional): For providing proxies, force_download, resume_download, cache_dir and other options specific to the 'from_pretrained' implementation where this will be supplied.
""" # noqa: ignore flake8"
if not config:
# if not ((encoder_name and decoder_name) or encoder_decoder_name) and not encoder_type:
if not ((encoder_name and decoder_name) or encoder_decoder_name):
raise ValueError(
"You must specify a Seq2Seq config \t OR \t"
"encoder_type, encoder_name, and decoder_name OR \t \t"
"encoder_type and encoder_decoder_name"
)
elif not (encoder_type or encoder_decoder_type):
raise ValueError(
"You must specify a Seq2Seq config \t OR \t"
"encoder_type, encoder_name, and decoder_name \t OR \t"
"encoder_type and encoder_decoder_name"
)
self.args = self._load_model_args(encoder_decoder_name)
if isinstance(args, dict):
self.args.update_from_dict(args)
elif isinstance(args, Seq2SeqArgs):
self.args = args
if "sweep_config" in kwargs:
sweep_config = kwargs.pop("sweep_config")
sweep_values = sweep_config_to_sweep_values(sweep_config)
self.args.update_from_dict(sweep_values)
if self.args.manual_seed:
random.seed(self.args.manual_seed)
np.random.seed(self.args.manual_seed)
torch.manual_seed(self.args.manual_seed)
if self.args.n_gpu > 0:
torch.cuda.manual_seed_all(self.args.manual_seed)
if use_cuda:
if torch.cuda.is_available():
if cuda_device == -1:
self.device = torch.device("cuda")
else:
self.device = torch.device(f"cuda:{cuda_device}")
else:
raise ValueError(
"'use_cuda' set to True when cuda is unavailable."
"Make sure CUDA is available or set `use_cuda=False`."
)
else:
self.device = "cpu"
self.results = {}
if not use_cuda:
self.args.fp16 = False
# config = EncoderDecoderConfig.from_encoder_decoder_configs(config, config)
if encoder_decoder_type:
config_class, model_class, tokenizer_class = MODEL_CLASSES[encoder_decoder_type]
else:
config_class, model_class, tokenizer_class = MODEL_CLASSES[encoder_type]
if encoder_decoder_type in ["bart", "marian"]:
self.model = model_class.from_pretrained(encoder_decoder_name)
if encoder_decoder_type == "bart":
self.encoder_tokenizer = tokenizer_class.from_pretrained(encoder_decoder_name)
elif encoder_decoder_type == "marian":
if self.args.base_marian_model_name:
self.encoder_tokenizer = tokenizer_class.from_pretrained(self.args.base_marian_model_name)
else:
self.encoder_tokenizer = tokenizer_class.from_pretrained(encoder_decoder_name)
self.decoder_tokenizer = self.encoder_tokenizer
self.config = self.model.config
else:
if encoder_decoder_name:
# self.model = EncoderDecoderModel.from_pretrained(encoder_decoder_name)
self.model = EncoderDecoderModel.from_encoder_decoder_pretrained(
os.path.join(encoder_decoder_name, "encoder"), os.path.join(encoder_decoder_name, "decoder")
)
self.encoder_tokenizer = tokenizer_class.from_pretrained(os.path.join(encoder_decoder_name, "encoder"))
self.decoder_tokenizer = AutoTokenizer.from_pretrained(os.path.join(encoder_decoder_name, "decoder"))
else:
self.model = EncoderDecoderModel.from_encoder_decoder_pretrained(
encoder_name, decoder_name, config=config
)
self.encoder_tokenizer = tokenizer_class.from_pretrained(encoder_name)
self.decoder_tokenizer = AutoTokenizer.from_pretrained(decoder_name)
self.encoder_config = self.model.config.encoder
self.decoder_config = self.model.config.decoder
if additional_special_tokens_encoder is not None:
self.encoder_tokenizer.add_special_tokens(additional_special_tokens_encoder)
if additional_special_tokens_decoder is not None:
self.decoder_tokenizer.add_special_tokens(additional_special_tokens_decoder)
if self.args.wandb_project and not wandb_available:
warnings.warn("wandb_project specified but wandb is not available. Wandb disabled.")
self.args.wandb_project = None
# `model_name` could be provided in args
if self.args.model_name is None:
if encoder_decoder_name:
self.args.model_name = encoder_decoder_name
# # Checking if we are loading from a saved model or using a pre-trained model
# if not saved_model_args and encoder_decoder_type == "marian":
# Need to store base pre-trained model name to get the tokenizer when loading a saved model
self.args.base_marian_model_name = encoder_decoder_name
elif encoder_name and decoder_name:
self.args.model_name = encoder_name + "-" + decoder_name
else:
self.args.model_name = "encoder-decoder"
if encoder_decoder_type:
self.args.model_type = encoder_decoder_type
elif encoder_type:
self.args.model_type = encoder_type + "-bert"
else:
self.args.model_type = "encoder-decoder"
def __init__(
self,
pretrained_model=None,
additional_special_tokens_encoder=None,
additional_special_tokens_decoder=None,
model_config=None,
vocab_file=None,
args=None,
use_cuda=True,
cuda_device=-1,
**kwargs,
):
self.args = self._load_model_args()
if isinstance(args, dict):
self.args.update_from_dict(args)
elif isinstance(args, Seq2SeqArgs):
self.args = args
if "sweep_config" in kwargs:
self.is_sweeping = True
sweep_config = kwargs.pop("sweep_config")
sweep_values = sweep_config_to_sweep_values(sweep_config)
self.args.update_from_dict(sweep_values)
else:
self.is_sweeping = False
if self.args.manual_seed:
random.seed(self.args.manual_seed)
np.random.seed(self.args.manual_seed)
torch.manual_seed(self.args.manual_seed)
if self.args.n_gpu > 0:
torch.cuda.manual_seed_all(self.args.manual_seed)
if use_cuda:
if torch.cuda.is_available():
if cuda_device == -1:
self.device = torch.device("cuda")
else:
self.device = torch.device(f"cuda:{cuda_device}")
else:
raise ValueError(
"'use_cuda' set to True when cuda is unavailable."
"Make sure CUDA is available or set `use_cuda=False`.")
else:
self.device = "cpu"
self.results = {}
if not use_cuda:
self.args.fp16 = False
# BartConfig, BartForConditionalGeneration, BartTokenizer
# config = EncoderDecoderConfig.from_encoder_decoder_configs(config, config)
model_config = BartConfig.from_json_file(model_config)
if pretrained_model is None:
self.model = BartForConditionalGeneration(config=model_config)
self.encoder_tokenizer = BartTokenizer.from_pretrained(vocab_file)
else:
self.model = BartForConditionalGeneration.from_pretrained(
pretrained_model)
self.encoder_tokenizer = BartTokenizer.from_pretrained(vocab_file)
self.decoder_tokenizer = self.encoder_tokenizer
# special AST token
# additional_special_tokens_encoder = {'additional_special_tokens': ['Assertion', 'RegExp', 'Repetition', 'Quantifier', 'ClassRange', 'CharacterClass']}
# additional_special_tokens_decoder = {'additional_special_tokens': ['Assertion', 'RegExp', 'Repetition', 'Quantifier', 'ClassRange', 'CharacterClass']}
self.config = self.model.config
if additional_special_tokens_encoder is not None:
self.encoder_tokenizer.add_special_tokens(
additional_special_tokens_encoder)
if additional_special_tokens_decoder is not None:
self.decoder_tokenizer.add_special_tokens(
additional_special_tokens_decoder)
if self.args.wandb_project and not wandb_available:
warnings.warn(
"wandb_project specified but wandb is not available. Wandb disabled."
)
self.args.wandb_project = None
self.args.model_type = 'bart'
self.args.model_name = 'ExplainREGEX'
