diff --git a/README.md b/README.md index 0e1c270..0b4fe38 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,28 @@ Supported methods: 3. P-Tuning 4. Prompt Tuning +## Getting started + +```python +from transformers import AutoModelForSeq2SeqLM +from pet import get_pet_config,get_pet_model +model_name_or_path = "bigscience/mt0-large" +tokenizer_name_or_path = "bigscience/mt0-large" + +config = { + "pet_type":"LORA", + "task_type":"SEQ_2_SEQ_LM", + "r": 8, + "lora_alpha": 32, + "lora_dropout": 0.1 +} +pet_config = get_pet_config(config) + +model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path) +model = get_pet_model(model, pet_config) +model.print_trainable_parameters() +``` + ## Models support matrix ### Sequence Classification diff --git a/src/pet/mapping.py b/src/pet/mapping.py index 1f2780f..daeb650 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -18,6 +18,7 @@ PET_TYPE_TO_CONFIG_MAPPING = { TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { "t5": ["q", "v"], + "mt5": ["q", "v"], "bart": ["q_proj", "v_proj"], "gpt2": ["c_attn"], "bloom": ["query_key_value"],