diff options
| -rw-r--r-- | baselines/peft_baseline.py | 30 | ||||
| -rw-r--r-- | scripts/run_all_methods.py | 8 |
2 files changed, 36 insertions, 2 deletions
diff --git a/baselines/peft_baseline.py b/baselines/peft_baseline.py index 442ba60..246385f 100644 --- a/baselines/peft_baseline.py +++ b/baselines/peft_baseline.py @@ -1,4 +1,4 @@ -"""PEFT baselines: LoRA, Tiny LoRA, and VeRA. +"""PEFT baselines: LoRA, Tiny LoRA, VeRA, Prompt Tuning, Prefix Tuning. Per-user adaptation on K support examples, then standard generation. Uses a class-based API to avoid repeated model wrapping/unwrapping. @@ -11,7 +11,10 @@ Usage: """ import torch -from peft import LoraConfig, VeraConfig, get_peft_model, TaskType +from peft import ( + LoraConfig, VeraConfig, PromptTuningConfig, PrefixTuningConfig, + get_peft_model, TaskType, PromptTuningInit, +) TARGET_MODULES = ["q_proj", "v_proj"] @@ -43,6 +46,21 @@ def _make_vera_config(rank, target_modules=None): ) +def _make_prompt_tuning_config(num_virtual_tokens): + return PromptTuningConfig( + task_type=TaskType.CAUSAL_LM, + num_virtual_tokens=num_virtual_tokens, + prompt_tuning_init=PromptTuningInit.RANDOM, + ) + + +def _make_prefix_tuning_config(num_virtual_tokens): + return PrefixTuningConfig( + task_type=TaskType.CAUSAL_LM, + num_virtual_tokens=num_virtual_tokens, + ) + + def get_lora_config(rank=8): return _make_lora_config(rank=rank) @@ -55,6 +73,14 @@ def get_vera_config(rank=256): return _make_vera_config(rank=rank) +def get_prompt_tuning_config(num_tokens=10): + return _make_prompt_tuning_config(num_virtual_tokens=num_tokens) + + +def get_prefix_tuning_config(num_tokens=10): + return _make_prefix_tuning_config(num_virtual_tokens=num_tokens) + + class PEFTBaseline: """Manages a PEFT-wrapped model for repeated per-user adaptation.""" diff --git a/scripts/run_all_methods.py b/scripts/run_all_methods.py index c5eb523..1502ff3 100644 --- a/scripts/run_all_methods.py +++ b/scripts/run_all_methods.py @@ -27,6 +27,7 @@ from adapt.cache_hidden import cache_support_hidden_states from adapt.fit_theta import fit_theta from baselines.peft_baseline import ( PEFTBaseline, get_lora_config, get_tiny_lora_config, get_vera_config, + get_prompt_tuning_config, get_prefix_tuning_config, ) from baselines.bm25_top1 import bm25_select_top1 from baselines.dense_retrieval import DenseRetriever @@ -38,6 +39,8 @@ ALL_METHODS = [ 'base', 'uph', 'prompt_all_k', 'bm25_top1', 'dense_top1', 'profile_based', 'lora', 'tiny_lora', 'vera', + 'prompt_tuning_5', 'prompt_tuning_10', 'prompt_tuning_20', + 'prefix_tuning_5', 'prefix_tuning_10', ] @@ -97,6 +100,11 @@ class MethodRunner: 'lora': lambda *a: self._run_peft(*a, config=get_lora_config(rank=8), lr=1e-4, desc='LoRA r=8'), 'tiny_lora': lambda *a: self._run_peft(*a, config=get_tiny_lora_config(rank=1), lr=1e-4, desc='Tiny LoRA r=1'), 'vera': lambda *a: self._run_peft(*a, config=get_vera_config(rank=256), lr=1e-3, desc='VeRA r=256'), + 'prompt_tuning_5': lambda *a: self._run_peft(*a, config=get_prompt_tuning_config(5), lr=3e-1, desc='PromptTuning L=5'), + 'prompt_tuning_10': lambda *a: self._run_peft(*a, config=get_prompt_tuning_config(10), lr=3e-1, desc='PromptTuning L=10'), + 'prompt_tuning_20': lambda *a: self._run_peft(*a, config=get_prompt_tuning_config(20), lr=3e-1, desc='PromptTuning L=20'), + 'prefix_tuning_5': lambda *a: self._run_peft(*a, config=get_prefix_tuning_config(5), lr=1e-2, desc='PrefixTuning L=5'), + 'prefix_tuning_10': lambda *a: self._run_peft(*a, config=get_prefix_tuning_config(10), lr=1e-2, desc='PrefixTuning L=10'), } if method_name not in dispatch: |