def __init__(self, model_type="BERT", model_name="bert-base-uncased"):
self.adaptor = get_adaptor(model_type)
model = AutoModelForNextSentencePrediction.from_pretrained(model_name)
super().__init__(model_type, model_name, model)
self._pipeline = None
self._trainer = None
def __init__(self,
model_type="DISTILBERT",
model_name="distilbert-base-uncased",
num_labels: int = 2,
load_path: str = ""):
self.adaptor = get_adaptor(model_type)
config = AutoConfig.from_pretrained(model_name, num_labels=num_labels)
if load_path != "":
model = AutoModelForSequenceClassification.from_pretrained(
load_path, config=config)
else:
model = AutoModelForSequenceClassification.from_pretrained(
model_name, config=config)
super().__init__(model_type, model_name, model)
device_number = detect_cuda_device_number()
self._pipeline = TextClassificationPipeline(model=self.model,
tokenizer=self.tokenizer,
device=device_number)
self._trainer = TCTrainer(self.model, self.model_type, self.tokenizer,
self._device, self.logger)
def __init__(
self, model_type: str = "DISTILBERT", model_name: str = "distilbert-base-uncased"):
self.adaptor = get_adaptor(model_type)
model = AutoModelForMaskedLM.from_pretrained(model_name)
super().__init__(model_type, model_name, model)
device_number = detect_cuda_device_number()
self._pipeline = FillMaskPipeline(model=self.model, tokenizer=self.tokenizer, device=device_number)
self._trainer = WPTrainer(self.model, model_type, self.tokenizer, self._device, self.logger)
def __init__(self,
model_type="DISTILBERT",
model_name="distilbert-base-cased-distilled-squad"):
self.adaptor = get_adaptor(model_type)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
super().__init__(model_type, model_name, model)
device_number = detect_cuda_device_number()
self._pipeline = QuestionAnsweringPipeline(model=self.model,
tokenizer=self.tokenizer,
device=device_number)
self._trainer = QATrainer(self.model, model_type, self.tokenizer,
self._device, self.logger)
def __init__(self,
model_type: str = "BERT",
model_name: str = "dslim/bert-base-NER"):
self.adaptor = get_adaptor(model_type)
model = AutoModelForTokenClassification.from_pretrained(model_name)
super().__init__(model_type, model_name, model)
device_number = detect_cuda_device_number()
self._pipeline = TokenClassificationPipeline(model=self.model,
tokenizer=self.tokenizer,
device=device_number)
self._trainer = TOCTrainer(self.model, model_type, self.tokenizer,
self._device, self.logger)
def __init__(self,
model_type="BERT",
model_name="bert-base-multilingual-cased",
num_labels=3):
self.adaptor = get_adaptor(model_type)
config = AutoConfig.from_pretrained(model_name, num_labels=num_labels)
model = AutoModelForSequenceClassification.from_pretrained(
model_name, config=config)
super().__init__(model_type, model_name, model)
device_number = detect_cuda_device_number()
self._pipeline = TextClassificationPipeline(model=self.model,
tokenizer=self.tokenizer,
device=device_number)
self._trainer = ABSATrainer(self.model, self.model_type,
self.tokenizer, self._device, self.logger)
def __init__(self,
model_type="DISTILBERT",
model_name="distilbert-base-uncased",
num_labels=2):
self.adaptor = get_adaptor(model_type)
config = AutoConfig.from_pretrained(model_name, num_labels=num_labels)
model = self.adaptor.SequenceClassification.from_pretrained(
model_name, config=config)
tokenizer = self.adaptor.Tokenizer.from_pretrained(model_name)
super().__init__(model_type, model_name, model, tokenizer)
device_number = detect_cuda_device_number()
self._pipeline = TextClassificationPipeline(model=model,
tokenizer=tokenizer,
device=device_number)
self._trainer = TCTrainer(self._model, self.model_type,
self._tokenizer, self._device, self.logger)
def __init__(self,
model_type: str = "GPT2",
model_name: str = "gpt2",
load_path: str = ""):
self.adaptor = get_adaptor(model_type)
if load_path != "":
model = AutoModelForCausalLM.from_pretrained(load_path)
else:
model = AutoModelForCausalLM.from_pretrained(model_name)
super().__init__(model_type, model_name, model)
device_number = detect_cuda_device_number()
self._pipeline = TextGenerationPipeline(model=self.model,
tokenizer=self.tokenizer,
device=device_number)
self._trainer = GENTrainer(self.model, model_type, self.tokenizer,
self._device, self.logger)