Ajuste e implante um modelo resumidor usando os contêineres Hugging Face Amazon SageMaker trazendo seu próprio script

Houve muitos avanços recentes no domínio da PNL. Modelos pré-treinados e serviços de PNL totalmente gerenciados democratizaram o acesso e a adoção da PNL. Amazon Comprehend é um serviço totalmente gerenciado que pode executar tarefas de NLP, como reconhecimento de entidade personalizada, modelagem de tópicos, análise de sentimentos e muito mais para extrair insights de dados sem a necessidade de qualquer experiência anterior em ML.

No ano passado, a AWS anunciou um parceria de Abraçando o rosto para ajudar a trazer modelos de processamento de linguagem natural (NLP) para a produção mais rapidamente. Hugging Face é uma comunidade de IA de código aberto, focada em PNL. Sua biblioteca baseada em Python (transformadores) fornece ferramentas para usar facilmente arquiteturas populares de Transformer de última geração, como BERT, RoBERTa e GPT. Você pode aplicar esses modelos a uma variedade de tarefas de PNL, como classificação de texto, extração de informações e resposta a perguntas, entre outras.

Amazon Sage Maker é um serviço totalmente gerenciado que fornece aos desenvolvedores e cientistas de dados a capacidade de criar, treinar e implantar modelos de aprendizado de máquina (ML) rapidamente. O SageMaker elimina o trabalho pesado de cada etapa do processo de ML, facilitando o desenvolvimento de modelos de alta qualidade. O SDK do Python do SageMaker fornece APIs e contêineres de código aberto para treinar e implantar modelos no SageMaker, usando várias estruturas diferentes de ML e aprendizado profundo.

A integração do Hugging Face com o SageMaker permite que você crie modelos do Hugging Face em escala em seus próprios casos de uso específicos do domínio.

Neste post, mostramos um exemplo de como criar e implantar um resumidor de texto Hugging Face personalizado no SageMaker. Usamos Pegasus [1] para este propósito, o primeiro modelo baseado em Transformer especificamente pré-treinado em um objetivo adaptado para sumarização de texto abstrato. O BERT é pré-treinado em mascarar palavras aleatórias em uma frase; em contraste, durante o pré-treinamento do Pegasus, as frases são mascaradas de um documento de entrada. O modelo então gera as sentenças faltantes como uma única sequência de saída usando todas as sentenças desmascaradas como contexto, criando um resumo executivo do documento como resultado.

Graças à flexibilidade da biblioteca HuggingFace, você pode facilmente adaptar o código mostrado neste post para outros tipos de modelos de transformadores, como t5, BART e muito mais.

Carregue seu próprio conjunto de dados para ajustar um modelo Hugging Face

Para carregar um conjunto de dados personalizado de um arquivo CSV, usamos o load_dataset do pacote Transformers. Podemos aplicar tokenização ao conjunto de dados carregado usando o datasets.Dataset.map função. o map A função itera sobre o conjunto de dados carregado e aplica a função tokenize a cada exemplo. O conjunto de dados tokenizado pode então ser passado para o treinador para ajustar o modelo. Veja o seguinte código:

# Python
def tokenize(batch):
    tokenized_input = tokenizer(batch[args.input_column], padding='max_length', truncation=True, max_length=args.max_source)
    tokenized_target = tokenizer(batch[args.target_column], padding='max_length', truncation=True, max_length=args.max_target)
    tokenized_input['target'] = tokenized_target['input_ids']

    return tokenized_input
    

def load_and_tokenize_dataset(data_dir):
    for file in os.listdir(data_dir):
        dataset = load_dataset("csv", data_files=os.path.join(data_dir, file), split='train')
    tokenized_dataset = dataset.map(lambda batch: tokenize(batch), batched=True, batch_size=512)
    tokenized_dataset.set_format('numpy', columns=['input_ids', 'attention_mask', 'labels'])
    
    return tokenized_dataset

Crie seu script de treinamento para o estimador Hugging Face SageMaker

Conforme explicado no post AWS e Hugging Face colaboram para simplificar e acelerar a adoção de modelos de processamento de linguagem natural, treinar um modelo Hugging Face no SageMaker nunca foi tão fácil. Podemos fazer isso usando o estimador Hugging Face do SDK do SageMaker.

O trecho de código a seguir ajusta o Pegasus em nosso conjunto de dados. Você também pode encontrar muitos cadernos de amostra que orientam você no ajuste fino de diferentes tipos de modelos, disponíveis diretamente no repositório GitHub de transformadores. Para habilitar o treinamento distribuído, podemos usar o Biblioteca de Paralelismo de Dados no SageMaker, que foi incorporado à API do HuggingFace Trainer. Para habilitar o paralelismo de dados, precisamos definir o distribution parâmetro em nosso estimador Hugging Face.

# Python
from sagemaker.huggingface import HuggingFace
# configuration for running training on smdistributed Data Parallel
distribution = {'smdistributed':{'dataparallel':{ 'enabled': True }}}
huggingface_estimator = HuggingFace(entry_point='train.py',
                                    source_dir='code',
                                    base_job_name='huggingface-pegasus',
                                    instance_type= 'ml.g4dn.16xlarge',
                                    instance_count=1,
                                    transformers_version='4.6',
                                    pytorch_version='1.7',
                                    py_version='py36',
                                    output_path=output_path,
                                    role=role,
                                    hyperparameters = {
                                        'model_name': 'google/pegasus-xsum',
                                        'epoch': 10,
                                        'per_device_train_batch_size': 2
                                    },
                                    distribution=distribution)
huggingface_estimator.fit({'train': training_input_path, 'validation': validation_input_path, 'test': test_input_path})

O tamanho máximo do lote de treinamento que você pode configurar depende do tamanho do modelo e da memória da GPU da instância usada. Se o treinamento distribuído do SageMaker estiver ativado, o tamanho total do lote será a soma de cada lote distribuído em cada dispositivo/GPU. Se usarmos uma instância ml.g4dn.16xlarge com treinamento distribuído em vez de uma instância ml.g4dn.xlarge, teremos oito vezes (8 GPUs) mais memória do que uma instância ml.g4dn.xlarge (1 GPU). O tamanho do lote por dispositivo permanece o mesmo, mas oito dispositivos estão treinando em paralelo.

Como de costume com o SageMaker, criamos um train.py script para usar com o Modo Script e passar hiperparâmetros para treinamento. O trecho de código a seguir para Pegasus carrega o modelo e o treina usando o Transformers Trainer classe:

# Python
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    Seq2SeqTrainer,
    Seq2seqTrainingArguments
)

model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
    
training_args = Seq2seqTrainingArguments(
    output_dir=args.model_dir,
    num_train_epochs=args.epoch,
    per_device_train_batch_size=args.train_batch_size,
    per_device_eval_batch_size=args.eval_batch_size,
    warmup_steps=args.warmup_steps,
    weight_decay=args.weight_decay,
    logging_dir=f"{args.output_data_dir}/logs",
    logging_strategy='epoch',
    evaluation_strategy='epoch',
    saving_strategy='epoch',
    adafactor=True,
    do_train=True,
    do_eval=True,
    do_predict=True,
    save_total_limit = 3,
    load_best_model_at_end=True,
    metric_for_best_model='eval_loss'
    # With the goal to deploy the best checkpoint to production
    # it is important to set load_best_model_at_end=True,
    # this makes sure that the last model is saved at the root
    # of the model_dir” directory
)
    
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset['train'],
    eval_dataset=dataset['validation']
)

trainer.train()
trainer.save_model()

# Get rid of unused checkpoints inside the container to limit the model.tar.gz size
os.system(f"rm -rf {args.model_dir}/checkpoint-*/")

O código completo está disponível em GitHub.

Implante o modelo Hugging Face treinado no SageMaker

Nossos amigos da Hugging Face tornaram a inferência sobre os modelos do SageMaker for Transformers mais simples do que nunca, graças ao Kit de ferramentas de inferência de rosto de abraço SageMaker. Você pode implantar diretamente o modelo previamente treinado simplesmente configurando a variável de ambiente "HF_TASK":"summarization" (para instruções, consulte Modelos Pegasus), escolhendo Implantação, e depois escolhendo Amazon Sage Maker, sem precisar escrever um script de inferência.

No entanto, se você precisar de alguma maneira específica de gerar ou pós-processar previsões, por exemplo, gerando várias sugestões de resumo com base em uma lista de diferentes parâmetros de geração de texto, escrever seu próprio script de inferência pode ser útil e relativamente simples:

# Python
# inference.py script

import os
import json
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def model_fn(model_dir):
    # Create the model and tokenizer and load weights
    # from the previous training Job, passed here through "model_dir"
    # that is reflected in HuggingFaceModel "model_data"
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_dir).to(device).eval()
    
    model_dict = {'model':model, 'tokenizer':tokenizer}
    
    return model_dict
        

def predict_fn(input_data, model_dict):
    # Return predictions/generated summaries
    # using the loaded model and tokenizer on input_data
    text = input_data.pop('inputs')
    parameters_list = input_data.pop('parameters_list', None)
    
    tokenizer = model_dict['tokenizer']
    model = model_dict['model']

    # Parameters may or may not be passed    
    input_ids = tokenizer(text, truncation=True, padding='longest', return_tensors="pt").input_ids.to(device)
    
    if parameters_list:
        predictions = []
        for parameters in parameters_list:
            output = model.generate(input_ids, **parameters)
            predictions.append(tokenizer.batch_decode(output, skip_special_tokens=True))
    else:
        output = model.generate(input_ids)
        predictions = tokenizer.batch_decode(output, skip_special_tokens=True)
    
    return predictions
    
    
def input_fn(request_body, request_content_type):
    # Transform the input request to a dictionary
    request = json.loads(request_body)
    return request

Conforme mostrado no código anterior, esse script de inferência para HuggingFace no SageMaker precisa apenas das seguintes funções de modelo:

• model_fn () – Lê o conteúdo do que foi salvo no final do trabalho de treinamento dentro SM_MODEL_DIR, ou de um diretório de pesos de modelo existente salvo como um arquivo tar.gz em Serviço de armazenamento simples da Amazon (Amazônia S3). Ele é usado para carregar o modelo treinado e o tokenizer associado.
• input_fn () – Formata os dados recebidos de uma solicitação feita ao terminal.
• Predict_fn () – Chama a saída de model_fn() (o modelo e o tokenizer) para executar inferência na saída de input_fn() (os dados formatados).

Opcionalmente, você pode criar um output_fn() função para formatação de inferência, usando a saída de predict_fn(), que não demonstramos neste post.

Podemos então implantar o modelo Hugging Face treinado com seu script de inferência associado ao SageMaker usando o Abraçando o rosto modelo SageMaker classe:

# Python
from sagemaker.huggingface import HuggingFaceModel

model = HuggingFaceModel(model_data=huggingface_estimator.model_data,
                     role=role,
                     framework_version='1.7',
                     py_version='py36',
                     entry_point='inference.py',
                     source_dir='code')
                     
predictor = model.deploy(initial_instance_count=1,
                         instance_type='ml.g4dn.xlarge'
                         )

Testar o modelo implantado

Para esta demonstração, treinamos o modelo no Conjunto de dados de avaliações de roupas de comércio eletrônico feminino, que contém resenhas de artigos de vestuário (que consideramos como texto de entrada) e seus títulos associados (que consideramos como resumos). Depois que removemos os artigos com títulos ausentes, o conjunto de dados contém 19,675 revisões. O ajuste fino do modelo Pegasus em um conjunto de treinamento contendo 70% desses artigos para cinco épocas levou aproximadamente 3.5 horas em uma instância ml.p3.16xlarge.

Podemos então implantar o modelo e testá-lo com alguns dados de exemplo do conjunto de teste. O seguinte é um exemplo de revisão descrevendo um suéter:

# Python
Review Text
"I ordered this sweater in green in petite large. The color and knit is beautiful and the shoulders and body fit comfortably; however, the sleeves were very long for a petite. I roll them, and it looks okay but would have rather had a normal petite length sleeve."

Original Title
"Long sleeves"

Rating
3

Graças ao nosso script de inferência personalizado hospedado em um endpoint SageMaker, podemos gerar vários resumos para esta revisão com diferentes parâmetros de geração de texto. Por exemplo, podemos pedir ao endpoint para gerar um intervalo de resumos muito curtos a moderadamente longos especificando diferentes penalidades de comprimento (quanto menor a penalidade de comprimento, menor o resumo gerado). A seguir estão alguns exemplos de entrada de parâmetros e os resumos gerados por máquina subsequentes:

# Python
inputs = {
    "inputs":[
"I ordered this sweater in green in petite large. The color and knit is   beautiful and the shoulders and body fit comfortably; however, the sleeves were very long for a petite. I roll them, and it looks okay but would have rather had a normal petite length sleeve."
    ],

    "parameters_list":[
        {
            "length_penalty":2
        },
	{
            "length_penalty":1
        },
	{
            "length_penalty":0.6
        },
        {
            "length_penalty":0.4
        }
    ]

result = predictor.predict(inputs)
print(result)

[
    ["Beautiful color and knit but sleeves are very long for a petite"],
    ["Beautiful sweater, but sleeves are too long for a petite"],
    ["Cute, but sleeves are long"],
    ["Very long sleeves"]
]

Qual resumo você prefere? O primeiro título gerado captura todos os fatos importantes sobre a resenha, com um quarto do número de palavras. Em contraste, o último usa apenas três palavras (menos de 1/10 do comprimento da revisão original) para se concentrar na característica mais importante do suéter.

Conclusão

Você pode ajustar um resumidor de texto em seu conjunto de dados personalizado e implantá-lo em produção no SageMaker com este exemplo simples disponível em GitHub. Adicional cadernos de amostra para treinar e implantar modelos de rosto abraçado no SageMaker também estão disponíveis.

Como sempre, a AWS agradece o feedback. Por favor, envie quaisquer comentários ou perguntas.

Referências

[1] PEGASUS: Pré-treinamento com frases-gap extraídas para sumarização abstrativa

Sobre os autores

Victor Malesevic é engenheiro de aprendizado de máquina com AWS Professional Services, apaixonado por processamento de linguagem natural e MLOps. Ele trabalha com clientes para desenvolver e colocar em produção modelos desafiadores de aprendizado profundo na AWS. Em seu tempo livre, ele gosta de compartilhar um copo de vinho tinto e um pouco de queijo com os amigos.

Aamna Najmi é um cientista de dados com AWS Professional Services. Ela é apaixonada por ajudar os clientes a inovar com tecnologias de Big Data e Inteligência Artificial para extrair valor de negócios e insights de dados. Em seu tempo livre, ela gosta de jardinagem e viajar para novos lugares.