Prefix Tuning

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PrefixTuningConfig, get_peft_model, TaskType

# 베이스 모델 로드
model_name = "gpt2-medium"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# Prefix Tuning 설정
peft_config = PrefixTuningConfig(
    task_type=TaskType.CAUSAL_LM,
    num_virtual_tokens=20,        # prefix 길이
    prefix_projection=True,       # MLP로 prefix 임베딩 생성
    encoder_hidden_size=1024      # 인코더 히든 사이즈
)

# PEFT 모델 생성
model = get_peft_model(model, peft_config)

# 학습 가능한 파라미터 확인
model.print_trainable_parameters()
# trainable params: 983,040 || all params: 355,350,528 || trainable%: 0.2767

# 학습 데이터 준비
from datasets import Dataset

train_data = [
    {"text": "고객: 배송이 언제 오나요? 답변: 주문 후 2-3일 내 배송됩니다."},
    {"text": "고객: 반품하고 싶어요. 답변: 7일 이내 무료 반품 가능합니다."},
    {"text": "고객: 결제 방법은? 답변: 신용카드, 계좌이체, 간편결제 가능합니다."}
]
dataset = Dataset.from_list(train_data)

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length",
                    truncation=True, max_length=128)

tokenized_dataset = dataset.map(tokenize_function, batched=True)

# 학습
from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./prefix_tuning_output",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=3e-2,  # prefix tuning은 높은 LR 사용
    logging_steps=10
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset
)

trainer.train()

# 추론
model.eval()
input_text = "고객: 교환하려면 어떻게 하나요? 답변:"
inputs = tokenizer(input_text, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

# 모델 저장 (prefix 파라미터만 저장됨)
model.save_pretrained("./my_prefix_model")

# === 멀티 태스크 서빙 (prefix 교체) ===
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained("gpt2-medium")

# 태스크별 prefix 로드
customer_service_model = PeftModel.from_pretrained(base_model, "./prefix_customer_service")
marketing_model = PeftModel.from_pretrained(base_model, "./prefix_marketing")
legal_model = PeftModel.from_pretrained(base_model, "./prefix_legal")

# 추론 시 태스크에 따라 prefix만 교체
def generate_response(query: str, task: str) -> str:
    if task == "customer_service":
        model = customer_service_model
    elif task == "marketing":
        model = marketing_model
    elif task == "legal":
        model = legal_model
    else:
        model = base_model

    inputs = tokenizer(query, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=100)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# === Prefix Tuning vs LoRA 비교 ===
# Prefix Tuning:
#   - 장점: 입력 길이가 증가하므로 추론 시 해석 가능
#   - 단점: prefix가 실제 컨텍스트를 차지함
#   - 파라미터: ~0.1-0.5%

# LoRA:
#   - 장점: 입력 길이 그대로, 더 좋은 성능
#   - 단점: weight 수정이라 해석 어려움
#   - 파라미터: ~0.1-1%

# 결론: 대부분의 경우 LoRA가 더 좋은 선택

# === Prompt Tuning (Prefix Tuning의 단순화 버전) ===
# Prompt Tuning은 입력 임베딩에만 prefix 추가
# Prefix Tuning은 모든 레이어의 hidden state에 추가

from peft import PromptTuningConfig, PromptTuningInit

prompt_tuning_config = PromptTuningConfig(
    task_type=TaskType.CAUSAL_LM,
    num_virtual_tokens=20,
    prompt_tuning_init=PromptTuningInit.RANDOM  # 또는 TEXT
)

# TEXT 초기화: 자연어로 프롬프트 초기화
# RANDOM 초기화: 랜덤 벡터로 시작