[AI] LoRA 모델과 원본 모델을 합치기

1. 현재 이런식으로 LoRA 모델이 만들어져 있다.

이걸 사용하려면 항상 원본 모델에 얹어서 로드해야만 했는데...

2. LoRA 모델과 원본 모델 merge

from transformers import AutoModelForCausalLM
from peft import PeftModel, PeftConfig

# 1. LoRA 설정 정보 로드 (여기서 base 모델 경로도 알 수 있음)
peft_model_path = "./TinyLlama/TinyLlama-1.1B-Chat-v1.0-iphone"
peft_config = PeftConfig.from_pretrained(peft_model_path)

# 2. 원본(base) 모델 로딩
base_model = AutoModelForCausalLM.from_pretrained(peft_config.base_model_name_or_path)

# 3. LoRA 모델 로딩 및 병합
model_with_lora = PeftModel.from_pretrained(base_model, peft_model_path)
merged_model = model_with_lora.merge_and_unload()  # LoRA 가중치를 병합

# 4. 병합된 모델 저장 (base + LoRA 합쳐진 상태)
merged_model.save_pretrained("./TinyLlama-1.1B-Chat-v1.0-iphone-merged")

# 5. tokenizer도 같이 저장 (원본 모델 기준)
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
tokenizer.save_pretrained("./TinyLlama-1.1B-Chat-v1.0-iphone-merged")

2-1. merge 결과

합치면 이런 모양새가 된다.

3. merge 된 모델로 추론

from transformers import AutoModelForCausalLM, AutoTokenizer

# 병합된 모델 경로
merged_model_dir = "./TinyLlama-1.1B-Chat-v1.0-iphone-merged"

print(f"merged model '{merged_model_dir}' 로드 중...")
model = AutoModelForCausalLM.from_pretrained(merged_model_dir)
print("merged model 로드 완료.")

# --- 토크나이저 로드 ---
print(f"토크나이저 '{merged_model_dir}' 로드 중...")
# 추론 시에도 학습 시 사용했던 동일한 토크나이저를 로드해야 함
tokenizer = AutoTokenizer.from_pretrained(merged_model_dir)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
print("토크나이저 로드 완료.")

# --- 추론 예시 ---
print("\n--- 추론 시작 ---")

# 모델에게 질문을 던질 메시지 리스트를 정의
# assistant 역할의 content는 모델이 생성할 부분
messages = [
    {"role": "system", "content": "당신은 아이폰 전문가 챗봇입니다. 사용자의 질문에 답변해 주세요."},
    {"role": "user", "content": "화면 캡처하는 방법 알려주세요."},
    {"role": "assistant", "content": ""}
]

# 챗 템플릿을 사용하여 모델 입력 형식에 맞는 문자열 프롬프트를 생성
# add_generation_prompt=True는 모델이 답변 생성을 시작해야 함을 나타내는 토큰을 추가
input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
print(f"입력 프롬프트:\n---\n{input_text}\n---")

# 생성된 문자열 프롬프트를 토큰 ID로 변환하고, 모델이 있는 장치로 이동
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

# 텍스트 생성 파라미터 설정
generation_params = {
    "max_new_tokens": 512,  # 생성할 최대 토큰 수
    "top_p": 0.9,           # Top-p 샘플링
    "temperature": 0.1,     # 온도 조절
    "eos_token_id": tokenizer.eos_token_id, # 문장 종료 토큰 지정
    "pad_token_id": tokenizer.pad_token_id  # 패딩 토큰 지정
}

print("텍스트 생성 중...")
outputs = model.generate(**inputs, **generation_params)
print("텍스트 생성 완료.")

# 생성된 토큰 ID를 다시 사람이 읽을 수 있는 텍스트로 변환
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"생성된 텍스트:\n---\n{generated_text}\n---")

print("\n--- 추론 완료 ---")

3-1. 추론 결과 확인

당연하겠지만 같은 결과가 나온다.