Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3from peft import PeftModel, PeftConfig
4
5# -----------------------------
6# 1. Define PEFT model ID & Checkpoint (Epoch)
7# -----------------------------
8peft_model_id = "xxccho/margin_reg_baseline_code"
9
10# [ Checkpoints to Epochs Mapping ]
11# Epoch 1 : "checkpoint-246"
12# Epoch 2 : "checkpoint-492"
13# Epoch 3 : "checkpoint-738"
14# Epoch 4 : "checkpoint-984"
15# Epoch 5 : "checkpoint-1230"
16# Epoch 6 : "checkpoint-1476"
17# Epoch 7 : "checkpoint-1722"
18# Epoch 8 : "checkpoint-1968"
19# Epoch 9 : "checkpoint-2214"
20# Epoch 10 : "checkpoint-2460"
21
22# 예시: 5 Epoch 체크포인트를 사용하려면 "checkpoint-1230" 할당. None일 시 최종(10x) 모델 로드
23checkpoint = None
24
25# 2. Load the PEFT config
26config = PeftConfig.from_pretrained(peft_model_id, subfolder=checkpoint) if checkpoint else PeftConfig.from_pretrained(peft_model_id)
27
28# 3. Load tokenizer
29tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
30if tokenizer.pad_token is None:
31 tokenizer.pad_token = tokenizer.eos_token
32
33# 4. Load base model
34base_model = AutoModelForSequenceClassification.from_pretrained(
35 config.base_model_name_or_path,
36 num_labels=1,
37 torch_dtype=torch.bfloat16,
38 device_map="auto"
39)
40
41# 5. Apply LoRA adapter
42model = PeftModel.from_pretrained(base_model, peft_model_id, subfolder=checkpoint) if checkpoint else PeftModel.from_pretrained(base_model, peft_model_id)
43model.config.pad_token_id = tokenizer.pad_token_id
44model.eval()
45
46# Example Usage
47text = "User: Write a python code for calculating fibonacci sequence.\nAssistant: Here is the code..."
48inputs = tokenizer(text, return_tensors="pt").to(model.device)
49
50with torch.no_grad():
51 outputs = model(**inputs)
52 reward_score = outputs.logits.squeeze().item()
53
54print(f"[Code] Reward Score: {reward_score:.4f}")
551@misc{vonwerra2022trl,
2 title = {{TRL: Transformer Reinforcement Learning}},
3 author = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec},
4 year = 2020,
5 journal = {GitHub repository},
6 publisher = {GitHub},
7 howpublished = {\url{https://github.com/huggingface/trl}}
8}