Update README.md

README.md

@@ -23,11 +23,35 @@ import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from peft import PeftModel, PeftConfig
 
-# 1. Define the PEFT model ID
+# -----------------------------
+# 1. Define the PEFT model ID & Checkpoint (Epoch)
+# -----------------------------
 peft_model_id = "xxccho/margin_reg_baseline"
 
-# 2. Load the PEFT config
-config = PeftConfig.from_pretrained(peft_model_id)
+# [Optional] 특정 Epoch의 중간 체크포인트를 불러오고 싶을 때 아래 변수를 지정하세요.
+# 지정하지 않고 None으로 두면 레포지토리 최상단에 있는 마지막(최종) 학습 모델이 로드됩니다.
+#
+# [ Ckeckpoints to Epochs Mapping ]
+# Epoch 1  : "checkpoint-246"
+# Epoch 2  : "checkpoint-492"
+# Epoch 3  : "checkpoint-738"
+# Epoch 4  : "checkpoint-984"
+# Epoch 5  : "checkpoint-1230"
+# Epoch 6  : "checkpoint-1476"
+# Epoch 7  : "checkpoint-1722"
+# Epoch 8  : "checkpoint-1968"
+# Epoch 9  : "checkpoint-2214"
+# Epoch 10 : "checkpoint-2460"
+
+# 예시: 5 Epoch 체크포인트를 사용하려면 아래와 같이 변경하세요.
+# checkpoint = "checkpoint-1230"  
+checkpoint = None  # None이면 디폴트로 제일 마지막 저장 모델을 씁니다.
+
+# 2. Load the PEFT config (체크포인트 지정 여부에 따라 subfolder 적용)
+if checkpoint:
+    config = PeftConfig.from_pretrained(peft_model_id, subfolder=checkpoint)
+else:
+    config = PeftConfig.from_pretrained(peft_model_id)
 
 # 3. Load tokenizer from base model (safer)
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
@@ -44,8 +68,12 @@ base_model = AutoModelForSequenceClassification.from_pretrained(
     device_map="auto"
 )
 
-# 5. Apply LoRA adapter
-model = PeftModel.from_pretrained(base_model, peft_model_id)
+# 5. Apply LoRA adapter (체크포인트 지정 여부에 따라 subfolder 적용)
+if checkpoint:
+    model = PeftModel.from_pretrained(base_model, peft_model_id, subfolder=checkpoint)
+else:
+    model = PeftModel.from_pretrained(base_model, peft_model_id)
+
 model.config.pad_token_id = tokenizer.pad_token_id
 model.eval()