def create_model(model_class, model_type, model_name, num_labels, weight, args,
use_cuda, cuda_device, **kwargs):
if model_class == "ClassificationModel":
return ClassificationModel(model_type, model_name, num_labels, weight,
args, use_cuda, cuda_device, **kwargs)
elif model_class == "MultiLabelClassificationModel":
return MultiLabelClassificationModel(model_type, model_name,
num_labels, weight, args,
use_cuda, cuda_device, **kwargs)
elif model_class == "QuestionAnsweringModel":
return QuestionAnsweringModel(model_type, model_name, args, use_cuda,
cuda_device, **kwargs)
elif model_class == "NERModel":
return NERModel(model_type,
model_name,
args=args,
use_cuda=use_cuda,
cuda_device=cuda_device,
**kwargs)
elif model_class == "T5Model":
args = T5Args()
args.use_multiprocessed_decoding = False
return T5Model(model_type,
model_name,
args=args,
use_cuda=use_cuda,
cuda_device=cuda_device,
**kwargs)
else:
raise ValueError(
"{} is either invalid or not yet implemented.".format(model_class))
def __init__(self,
model_path: str = config.DEFAULT_MODEL_PATH,
model_architecture: str = config.MODEL_ARCHITECTURE,
use_cuda: bool = config.GPU):
'''
Constructs all the necessary attributes for the MT5_Translator object.
Parameters
----------
model_path : str
path to the mt5_translator model
model_architecture : str
model architecture (mt5, t5 ...)
use_cuda : bool
whether to use CUDA or not (if available)
'''
logging.info("Loading model...")
self.model_path = model_path
self.use_cuda = use_cuda
self.device = True \
if torch.cuda.is_available() and self.use_cuda else False
self.model_args = T5Args()
self.model_args.max_length = 512
self.model_args.length_penalty = 1
self.model_args.num_beams = 10
self.model = T5Model("mt5",
self.model_path,
args=self.model_args,
use_cuda=self.device)
logging.info(f"Use CUDA: {self.device}")
logging.info(f"Num GPUs Available: {torch.cuda.device_count()}")
logging.info(f"Model loaded")
def translate(self,
sentences: List[str],
max_length: int = 512,
min_length: int = None,
num_beams: int = 4):
'''
Generate translations from sentences
Parameters
----------
sentences : list of str
list of sentences. Each sentence gets its own translation
'''
self.model_args = T5Args()
self.model_args.max_length = max_length if max_length else max_length
self.model_args.min_length = min_length if min_length else min_length
self.model_args.num_beams = num_beams if num_beams else num_beams
self.model.args = self.model_args
return self.model.predict(sentences)
import logging
import pandas as pd
from simpletransformers.t5 import T5Model, T5Args
logging.basicConfig(level=logging.INFO)
transformers_logger = logging.getLogger("transformers")
transformers_logger.setLevel(logging.WARNING)
train_df = pd.read_csv("data/eng-spa/train.tsv", sep="\t").astype(str)
eval_df = pd.read_csv("data/eng-spa/eval.tsv", sep="\t").astype(str)
train_df["prefix"] = ""
eval_df["prefix"] = ""
model_args = T5Args()
model_args.max_seq_length = 80
model_args.train_batch_size = 10
model_args.eval_batch_size = 10
model_args.num_train_epochs = 1
model_args.evaluate_during_training = True
model_args.evaluate_during_training_steps = 30000
model_args.use_multiprocessing = False
model_args.fp16 = False
model_args.save_steps = -1
model_args.save_eval_checkpoints = False
model_args.no_cache = True
model_args.reprocess_input_data = True
model_args.overwrite_output_dir = True
model_args.preprocess_inputs = False
model_args.num_return_sequences = 1