Package: alexa-teacher-models
Alexa Teacher ModelsThis is the official Alexa Teacher Model program github page.AlexaTM 20BAlexaTM 20B is a 20B-Parameter sequence-to-sequence transformer model created by the Alexa Teacher Model (AlexaTM) team at Amazon. The model was trained on a mixture of Common Crawl (mC4) and Wikipedia data across 12 languages using denoising and Causal Language Modeling (CLM) tasks.AlexaTM 20B can be used for in-context learning. "In-context learning," also known as "prompting," refers to a method for using NLP models in which no fine tuning is required per task. Training examples are provided to the model only as part of the prompt given as inference input, a paradigm known as "few-shot in-context learning." In some cases, the model can perform well without any training data at all, a paradigm known as "zero-shot in-context learning."To learn more about the model, please read theAmazon Science blog postand thepaper.The model is currently available for noncommercial use via SageMaker JumpStart, as described in ourAWS blog post. The model can be accessed using the following steps:Createan AWS account if needed.In your AWS account, search forSageMakerin the search bar and click on it.Once in the SageMaker experience, create adomainand a studio user if none yet exist. All of the default settings can be used.In the control panel, clickLaunch appnext to the user you wish to use. Launch a studio instance.Once in the studio, there will be a launcher showing JumpStart as one of the tiles. ClickGo to SageMaker Jumpstart. Alternatively, JumpStart can be accessed by 3-pointed orange symbol on the far left of the studio.Once in JumpStart, click theNotebooksbutton.Browse or search for our example notebook entitledIn-context learning with AlexaTM 20B.There will be a button at the top to copy the read-only version into your studio.Ensure that your kernel has started, and run the notebook.Note: You can also find our example notebookhereLoad the Model and Run Inferencefromalexa_teacher_modelsimportAlexaTMTokenizerFasttokenizer=AlexaTMTokenizerFast.from_pretrained('/path/to/AlexaTM-20B-pr/')# Load the modelfromalexa_teacher_modelsimportAlexaTMSeq2SeqForConditionalGenerationmodel=AlexaTMSeq2SeqForConditionalGeneration.from_pretrained('/path/to/AlexaTM-20B-pr/')You can also use theAutoTokenizerandAutoModelForSeq2SeqLMas you would in any other HuggingFace Transformer
program by importingalexa_teacher_models:importalexa_teacher_models...tokenizer=AutoTokenizer.from_pretrained('/path/to/AlexaTM-20B-pr/')model=AutoModelForSeq2SeqLM.from_pretrained('/path/to/AlexaTM-20B-pr/')Load the model on 4 gpus:model.bfloat16()model.parallelize(4)Run the model in CLM mode:# qatest="""[CLM] Question: Who is the vocalist of coldplay? Answer:"""print('Input:',test)encoded=tokenizer(test,return_tensors="pt").to('cuda:0')generated_tokens=model.generate(input_ids=encoded['input_ids'],max_length=32,num_beams=1,num_return_sequences=1,early_stopping=True)tokenizer.batch_decode(generated_tokens,skip_special_tokens=True)[0]Run the model in denoising mode:# denoisingtest="we went to which is the capital of France"print('Input:',test)encoded=tokenizer(test,return_tensors="pt").to('cuda:0')generated_tokens=model.generate(input_ids=encoded['input_ids'],max_length=32,num_beams=5,num_return_sequences=5,early_stopping=True)tokenizer.batch_decode(generated_tokens,skip_special_tokens=True)Running the repl exampleA sample Read Execute Print Loop (REPL) program is provided in the samples.  It can be used to interact with
any AlexaTM model, and has a flexible set of command line arguments, including support for sampling and using multiple turns of history as context$ pip install alexa_teacher_models[repl]
$ python -m alexa_teacher_models.scripts.repl --model /path/to/AlexaTM-20B-pr/ --max_length 64
$ python -m alexa_teacher_models.scripts.repl --model /path/to/AlexaTM-20B-pr/ --max_length 64 --do_sample --max_history 3 --join_string " 
 "Fine-tuning with DeepSpeed on a single P4NoteWe strongly recommend training on multiple instances.  For information on how to do this, see the section belowTo run on a single P4 (8 GPUs), you will need to use CPU offload.  A deepspeed config is provided in thescripts/deepspeeddirectory.
Assuming you have a training and validation JSONL formatted file, a run would look like this:$ pip install alexa_teacher_models[ft]
$ deepspeed --num_gpus 8 --module alexa_teacher_models.scripts.finetune --per_device_train_batch_size $BS \
    --deepspeed deepspeed/zero3-offload.json \
    --model_name_or_path /home/ubuntu/AlexaTM/ --max_length 512 --bf16 --output_dir output \
    --max_target_length 64 --do_train --learning_rate 1e-7 \
    --train_file train.json --validation_file valid.json \
    --num_train_epochs 1 --save_steps 1000Fine-tuning with DeepSpeed on multiple machinesThere is adetailed tutorialdemonstrating how to fine-tune 20B across multiple machines in EC2 usingElastic Fabric Adapter (EFA).CitationIf you use AlexaTM 20B, please use the following BibTeX entry.@article{soltan2022alexatm,
  title={AlexaTM 20B: Few-Shot Learning Using a Large-Scale Multilingual Seq2seq Model},
  author={Saleh Soltan, Shankar Ananthakrishnan, Jack FitzGerald, Rahul Gupta, Wael Hamza, Haidar Khan, Charith Peris, Stephen Rawls, Andy Rosenbaum, Anna Rumshisky, Chandana Satya Prakash, Mukund Sridhar, Fabian Triefenbach, Apurv Verma, Gokhan Tur, Prem Natarajan},
  year={2022}
}SecuritySeeCONTRIBUTINGfor more information.LicenseThe code in this package is subject toLicense. However,
the model weights are subject toModel License.