| 40 | |
| 41 | |
| 42 | def create_peft_model(model, args): |
| 43 | |
| 44 | if 'roberta' in args.model: |
| 45 | |
| 46 | peft_config = LoraConfig( |
| 47 | task_type=TaskType.SEQ_CLS, |
| 48 | r=args.lora_r, |
| 49 | lora_alpha=args.lora_alpha, |
| 50 | lora_dropout=args.lora_dropout, |
| 51 | target_modules=["query", "value", "attention.output.dense", "output.dense"], |
| 52 | ) |
| 53 | |
| 54 | elif 't5' in args.model: |
| 55 | |
| 56 | peft_config = LoraConfig( |
| 57 | task_type=TaskType.SEQ_2_SEQ_LM, |
| 58 | r=args.lora_r, |
| 59 | lora_alpha=args.lora_alpha, |
| 60 | lora_dropout=args.lora_dropout, |
| 61 | target_modules=["q", "v", "k", "o", "wi", "wo"], |
| 62 | ) |
| 63 | |
| 64 | model = get_peft_model(model, peft_config) |
| 65 | |
| 66 | model.to(args.device) |
| 67 | |
| 68 | return model, peft_config |
| 69 | |
| 70 | |
| 71 | def create_model_tokenizer_it(args): |