Finjustera och distribuera en sammanfattningsmodell med hjälp av Hugging Face Amazon SageMaker-behållare som tar med ditt eget manus

Det har gjorts många nya framsteg inom NLP-domänen. Förutbildade modeller och helt hanterade NLP-tjänster har demokratiserat åtkomst och adoption av NLP. Amazon Comprehend är en helt hanterad tjänst som kan utföra NLP-uppgifter som anpassad enhetsigenkänning, ämnesmodellering, sentimentanalys och mer för att extrahera insikter från data utan att behöva ha någon tidigare ML-erfarenhet.

Förra året tillkännagav AWS en partnerskap med Kramande ansikte för att hjälpa till att producera NLP-modeller (natural language processing) snabbare. Hugging Face är en AI-gemenskap med öppen källkod, fokuserad på NLP. Deras Python-baserade bibliotek (Transformatorer) tillhandahåller verktyg för att enkelt använda populära toppmoderna transformatorarkitekturer som BERT, RobERTa och GPT. Du kan tillämpa dessa modeller på en mängd olika NLP-uppgifter, såsom textklassificering, informationsextraktion och frågesvar, bland annat andra.

Amazon SageMaker är en helt hanterad tjänst som ger utvecklare och datavetare möjligheten att snabbt bygga, träna och distribuera modeller för maskininlärning (ML). SageMaker tar bort de tunga lyften från varje steg i ML-processen, vilket gör det lättare att utveckla högkvalitativa modeller. SageMaker Python SDK tillhandahåller API:er och behållare med öppen källkod för att träna och distribuera modeller på SageMaker, med hjälp av flera olika ramverk för ML och djupinlärning.

Hugging Face-integrationen med SageMaker låter dig bygga Hugging Face-modeller i stor skala på dina egna domänspecifika användningsfall.

I det här inlägget går vi igenom ett exempel på hur man bygger och distribuerar en anpassad Hugging Face-textsammanfattning på SageMaker. Vi använder Pegasus [1] för detta ändamål, den första transformatorbaserade modellen som är specifikt förtränad på ett mål skräddarsytt för abstrakt textsammanfattning. BERT är förtränad i att maskera slumpmässiga ord i en mening; däremot, under Pegasus förträning, maskeras meningar från ett inmatningsdokument. Modellen genererar sedan de saknade meningarna som en enda utdatasekvens med alla omaskerade meningar som sammanhang, vilket skapar en sammanfattning av dokumentet som ett resultat.

Tack vare HuggingFace-bibliotekets flexibilitet kan du enkelt anpassa koden som visas i det här inlägget för andra typer av transformatormodeller, som t5, BART och mer.

Ladda din egen datauppsättning för att finjustera en Hugging Face-modell

För att ladda en anpassad datauppsättning från en CSV-fil använder vi load_dataset metod från Transformers-paketet. Vi kan tillämpa tokenisering på den laddade datamängden med hjälp av datasets.Dataset.map fungera. De map funktion itererar över den laddade datamängden och tillämpar tokenize-funktionen på varje exempel. Den tokeniserade datamängden kan sedan skickas till tränaren för att finjustera modellen. Se följande kod:

# 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

Bygg ditt träningsmanus för Hugging Face SageMaker-skattaren

Som förklarat i inlägget AWS och Hugging Face samarbetar för att förenkla och påskynda användningen av Natural Language Processing-modeller, att träna en Hugging Face-modell på SageMaker har aldrig varit enklare. Vi kan göra det genom att använda estimatorn för Hugging Face från SageMaker SDK.

Följande kodavsnitt finjusterar Pegasus på vår datauppsättning. Du kan också hitta många exempel anteckningsböcker som guidar dig genom att finjustera olika typer av modeller, tillgängliga direkt i transformatorernas GitHub-repository. För att möjliggöra distribuerad utbildning kan vi använda Data Parallelism Library i SageMaker, som har byggts in i HuggingFace Trainer API. För att möjliggöra dataparallellism måste vi definiera distribution parameter i vår Hugging Face-estimator.

# 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})

Den maximala träningsbatchstorleken du kan konfigurera beror på modellstorleken och GPU-minnet för den instans som används. Om SageMaker distribuerad utbildning är aktiverad är den totala batchstorleken summan av varje batch som distribueras över varje enhet/GPU. Om vi ​​använder en ml.g4dn.16xlarge med distribuerad träning istället för en ml.g4dn.xlarge-instans har vi åtta gånger (8 GPU) så mycket minne som en ml.g4dn.xlarge-instans (1 GPU). Batchstorleken per enhet förblir densamma, men åtta enheter tränas parallellt.

Som vanligt med SageMaker skapar vi en train.py skript att använda med skriptläge och skicka hyperparametrar för träning. Följande kodsnutt för Pegasus laddar modellen och tränar den med hjälp av Transformers Trainer klass:

# 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-*/")

Hela koden finns tillgänglig på GitHub.

Distribuera den utbildade Hugging Face-modellen till SageMaker

Våra vänner på Hugging Face har gjort slutsatser om SageMaker for Transformers-modeller enklare än någonsin tack vare SageMaker Hugging Face Inference Toolkit. Du kan distribuera den tidigare tränade modellen direkt genom att helt enkelt ställa in miljövariabeln "HF_TASK":"summarization" (för instruktioner, se Pegasus-modeller), väljer Distribuera, och sedan välja Amazon SageMaker, utan att behöva skriva ett slutledningsskript.

Men om du behöver något specifikt sätt att generera eller efterbehandla förutsägelser, till exempel generera flera sammanfattningsförslag baserat på en lista med olika textgenereringsparametrar, kan det vara användbart och relativt enkelt att skriva ditt eget slutledningsskript:

# 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

Som visas i föregående kod behöver ett sådant slutledningsskript för HuggingFace på SageMaker bara följande mallfunktioner:

• model_fn () – Läser innehållet i det som sparades i slutet av träningsjobbet inuti SM_MODEL_DIR, eller från en befintlig modellviktkatalog sparad som en tar.gz-fil i Amazon enkel lagringstjänst (Amazon S3). Den används för att ladda den tränade modellen och tillhörande tokenizer.
• input_fn () – Formaterar data som tagits emot från en begäran till slutpunkten.
• predict_fn () – Anropar utgången för model_fn() (modellen och tokenizern) för att göra slutledningar om utmatningen av input_fn() (de formaterade uppgifterna).

Alternativt kan du skapa en output_fn() funktion för slutledningsformatering, med utdata från predict_fn(), vilket vi inte visade i det här inlägget.

Vi kan sedan distribuera den tränade Hugging Face-modellen med tillhörande slutledningsskript till SageMaker med hjälp av Kramande ansikte SageMaker-modell klass:

# 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'
                         )

Testa den utplacerade modellen

För den här demon tränade vi modellen på Datauppsättning av recensioner av kläder för e-handel för kvinnor, som innehåller recensioner av klädartiklar (som vi betraktar som ingångstext) och deras tillhörande titlar (som vi betraktar som sammanfattningar). Efter att vi tagit bort artiklar med saknade titlar innehåller datasetet 19,675 70 recensioner. Att finjustera Pegasus-modellen på ett träningsset som innehåller 3.5 % av dessa artiklar under fem epoker tog cirka 3.16 timmar på en ml.pXNUMXxlarge instans.

Vi kan sedan distribuera modellen och testa den med några exempeldata från testsetet. Följande är ett exempel på en recension som beskriver en tröja:

# 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

Tack vare vårt anpassade slutledningsskript som finns i en SageMaker-slutpunkt kan vi generera flera sammanfattningar för denna recension med olika textgenereringsparametrar. Till exempel kan vi be slutpunkten att generera en rad mycket korta till måttligt långa sammanfattningar som anger olika längder på straff (ju mindre strafflängd, desto kortare genererad sammanfattning). Följande är några exempel på parameterinmatning och de efterföljande maskingenererade sammanfattningarna:

# 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"]
]

Vilken sammanfattning föredrar du? Den första genererade titeln fångar alla viktiga fakta om recensionen, med en fjärdedel av antalet ord. Däremot använder den sista bara tre ord (mindre än 1/10 av längden på den ursprungliga recensionen) för att fokusera på tröjans viktigaste egenskap.

Slutsats

Du kan finjustera en textsammanfattning på din anpassade datauppsättning och distribuera den till produktion på SageMaker med detta enkla exempel tillgängligt på GitHub. Ytterligare exempel anteckningsböcker att träna och distribuera Hugging Face-modeller på SageMaker är också tillgängliga.

Som alltid välkomnar AWS feedback. Skicka eventuella kommentarer eller frågor.

Referensprojekt

[1] PEGASUS: Förträning med extraherade gap-satser för abstrakt sammanfattning

Om författarna

Viktor Malesevic är en maskininlärningsingenjör med AWS Professional Services, passionerad för Natural Language Processing och MLOps. Han arbetar med kunder för att utveckla och sätta utmanande modeller för djupinlärning i produktion på AWS. På fritiden delar han gärna ett glas rött vin och lite ost med vänner.

Aamna Najmi är en dataforskare med AWS Professional Services. Hon brinner för att hjälpa kunder att förnya sig med Big Data och Artificiell Intelligens-teknik för att utnyttja affärsvärde och insikter från data. På fritiden tycker hon om trädgårdsarbete och att resa till nya platser.