transfer-llm 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (31) hide show
  1. transfer_llm-0.1.0/LICENSE +21 -0
  2. transfer_llm-0.1.0/PKG-INFO +247 -0
  3. transfer_llm-0.1.0/README.md +217 -0
  4. transfer_llm-0.1.0/pyproject.toml +58 -0
  5. transfer_llm-0.1.0/setup.cfg +4 -0
  6. transfer_llm-0.1.0/test/__init__.py +0 -0
  7. transfer_llm-0.1.0/test/unit/__init__.py +0 -0
  8. transfer_llm-0.1.0/test/unit/evaluate/__init__.py +0 -0
  9. transfer_llm-0.1.0/test/unit/evaluate/test_perplex.py +141 -0
  10. transfer_llm-0.1.0/test/unit/evaluate/test_perplex_output.py +100 -0
  11. transfer_llm-0.1.0/tests/test_cli.py +174 -0
  12. transfer_llm-0.1.0/tests/test_config.py +115 -0
  13. transfer_llm-0.1.0/tests/test_dpo_strategy.py +189 -0
  14. transfer_llm-0.1.0/tests/test_sft_strategy.py +282 -0
  15. transfer_llm-0.1.0/tests/test_trainer.py +358 -0
  16. transfer_llm-0.1.0/transfer/__init__.py +22 -0
  17. transfer_llm-0.1.0/transfer/cli.py +335 -0
  18. transfer_llm-0.1.0/transfer/config.py +139 -0
  19. transfer_llm-0.1.0/transfer/evaluation.py +371 -0
  20. transfer_llm-0.1.0/transfer/strategies/__init__.py +0 -0
  21. transfer_llm-0.1.0/transfer/strategies/base.py +26 -0
  22. transfer_llm-0.1.0/transfer/strategies/dpo.py +118 -0
  23. transfer_llm-0.1.0/transfer/strategies/sft.py +244 -0
  24. transfer_llm-0.1.0/transfer/trainer.py +648 -0
  25. transfer_llm-0.1.0/transfer/utils.py +14 -0
  26. transfer_llm-0.1.0/transfer_llm.egg-info/PKG-INFO +247 -0
  27. transfer_llm-0.1.0/transfer_llm.egg-info/SOURCES.txt +29 -0
  28. transfer_llm-0.1.0/transfer_llm.egg-info/dependency_links.txt +1 -0
  29. transfer_llm-0.1.0/transfer_llm.egg-info/entry_points.txt +2 -0
  30. transfer_llm-0.1.0/transfer_llm.egg-info/requires.txt +16 -0
  31. transfer_llm-0.1.0/transfer_llm.egg-info/top_level.txt +1 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2024 Dedy
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,247 @@
1
+ Metadata-Version: 2.4
2
+ Name: transfer-llm
3
+ Version: 0.1.0
4
+ Summary: A modular PyTorch framework for fine-tuning Hugging Face models.
5
+ Author-email: Dedy <dedy.ariansyah@gmail.com>
6
+ Project-URL: Homepage, https://github.com/deduu/transfer
7
+ Project-URL: Issues, https://github.com/deduu/transfer/issues
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Operating System :: OS Independent
11
+ Requires-Python: >=3.11
12
+ Description-Content-Type: text/markdown
13
+ License-File: LICENSE
14
+ Requires-Dist: numpy>=1.21.6
15
+ Requires-Dist: accelerate>=1.0.0
16
+ Requires-Dist: bitsandbytes>=0.44.0
17
+ Requires-Dist: transformers
18
+ Requires-Dist: datasets
19
+ Requires-Dist: tqdm
20
+ Requires-Dist: pyyaml
21
+ Requires-Dist: peft
22
+ Requires-Dist: sentence-transformers
23
+ Requires-Dist: scikit-learn
24
+ Requires-Dist: scipy
25
+ Requires-Dist: hf-xet>=1.2.0
26
+ Requires-Dist: torch>=2.3.0
27
+ Provides-Extra: dev
28
+ Requires-Dist: pytest>=7.0; extra == "dev"
29
+ Dynamic: license-file
30
+
31
+ # Transfer
32
+
33
+ A modular PyTorch framework for fine-tuning Hugging Face language models.
34
+
35
+ Transfer simplifies applying techniques like Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO) to any Hugging Face causal language model.
36
+
37
+ ## Features
38
+
39
+ - **SFT and DPO** — two fine-tuning strategies out of the box
40
+ - **Multi-turn conversation support** — point to a `messages_column` of chat-style message lists
41
+ - **Completion-only loss masking** — train on assistant tokens only with `train_on_completions_only`
42
+ - **LoRA / QLoRA** — 4-bit quantization (nf4 / fp4) via bitsandbytes, configurable LoRA rank and target modules
43
+ - **Training loop controls** — gradient accumulation, LR scheduling (cosine, linear, constant, …), gradient clipping
44
+ - **Checkpointing and logging** — intermediate checkpoints every N steps, CSV training logs, optional W&B integration
45
+ - **Built-in evaluation** — perplexity and semantic entropy metrics
46
+ - **CLI** — `transfer train` and `transfer infer` commands for training and inference without writing Python
47
+
48
+ ## Installation
49
+
50
+ ### From PyPI
51
+
52
+ ```bash
53
+ pip install transfer-llm
54
+ ```
55
+
56
+ ### From source
57
+
58
+ ```bash
59
+ git clone https://github.com/deduu/transfer.git
60
+ cd transfer
61
+ pip install .
62
+ ```
63
+
64
+ ## Quick Start
65
+
66
+ ### Single-turn SFT
67
+
68
+ ```python
69
+ from datasets import Dataset
70
+ from transfer import Trainer, SFTConfig
71
+
72
+ # Prepare data
73
+ dataset = Dataset.from_dict({
74
+ "prompt": ["What is the capital of France?", "Explain gravity."],
75
+ "response": ["The capital of France is Paris.", "Gravity is ..."],
76
+ })
77
+
78
+ # Configure
79
+ config = SFTConfig(
80
+ model_name="google/gemma-2b",
81
+ num_epochs=3,
82
+ batch_size=2,
83
+ learning_rate=5e-5,
84
+ output_dir="./gemma-sft",
85
+ )
86
+
87
+ # Train
88
+ trainer = Trainer(task="sft", config=config, train_dataset=dataset)
89
+ trainer.train()
90
+ trainer.save_model()
91
+ ```
92
+
93
+ ### Multi-turn SFT with completion masking and LoRA
94
+
95
+ ```python
96
+ from datasets import Dataset
97
+ from transfer import Trainer, SFTConfig
98
+
99
+ # Each row contains a list of message dicts
100
+ conversations = [
101
+ {"messages": [
102
+ {"role": "system", "content": "You are a helpful assistant."},
103
+ {"role": "user", "content": "What is the weather in SF?"},
104
+ {"role": "assistant", "content": "It is 62 F and foggy."},
105
+ ]},
106
+ # ... more conversations
107
+ ]
108
+ dataset = Dataset.from_list(conversations)
109
+
110
+ config = SFTConfig(
111
+ model_name="meta-llama/Llama-3.2-3B-Instruct",
112
+ messages_column="messages", # read from chat-style column
113
+ train_on_completions_only=True, # loss on assistant tokens only
114
+ num_epochs=3,
115
+ batch_size=1,
116
+ learning_rate=2e-5,
117
+ max_length=1024,
118
+ gradient_accumulation_steps=4,
119
+ warmup_steps=10,
120
+ scheduler_type="cosine",
121
+ save_steps=50,
122
+ logging_steps=5,
123
+ use_lora=True,
124
+ lora_r=16,
125
+ lora_alpha=32,
126
+ lora_dropout=0.05,
127
+ output_dir="./agent-sft-output",
128
+ )
129
+
130
+ trainer = Trainer(task="sft", config=config, train_dataset=dataset)
131
+ trainer.train()
132
+ trainer.save_model()
133
+ ```
134
+
135
+ ### DPO
136
+
137
+ ```python
138
+ from datasets import load_dataset
139
+ from transfer import Trainer, DPOConfig
140
+
141
+ dataset = load_dataset("Anthropic/hh-rlhf", split="train[:1%]")
142
+
143
+ config = DPOConfig(
144
+ model_name="google/gemma-2b",
145
+ num_epochs=1,
146
+ batch_size=2,
147
+ beta=0.1,
148
+ output_dir="./gemma-dpo",
149
+ )
150
+
151
+ trainer = Trainer(task="dpo", config=config, train_dataset=dataset)
152
+ trainer.train()
153
+ trainer.save_model()
154
+ ```
155
+
156
+ ## CLI Usage
157
+
158
+ The `transfer` command is installed automatically with the package.
159
+
160
+ ### Training
161
+
162
+ ```bash
163
+ transfer train \
164
+ --task sft \
165
+ --model_name google/gemma-2b \
166
+ --dataset_path data.jsonl \
167
+ --output_dir ./my-sft-model \
168
+ --num_epochs 3 \
169
+ --batch_size 4 \
170
+ --learning_rate 2e-4 \
171
+ --use_lora \
172
+ --gradient_accumulation_steps 4 \
173
+ --scheduler_type cosine \
174
+ --warmup_steps 50 \
175
+ --save_steps 200 \
176
+ --logging_steps 10
177
+ ```
178
+
179
+ Use `--messages_column messages --train_on_completions_only` for multi-turn datasets.
180
+
181
+ ### Inference
182
+
183
+ ```bash
184
+ transfer infer \
185
+ --model_name google/gemma-2b \
186
+ --adapter_path ./my-sft-model \
187
+ --prompt "Explain quantum computing" \
188
+ --use_chat_template \
189
+ --max_new_tokens 256 \
190
+ --temperature 0.7 \
191
+ --do_sample
192
+ ```
193
+
194
+ ## Configuration
195
+
196
+ ### SFTConfig
197
+
198
+ | Parameter | Default | Description |
199
+ |---|---|---|
200
+ | `model_name` | `"google/gemma-2b"` | Hugging Face model name or path |
201
+ | `num_epochs` | `30` | Number of training epochs |
202
+ | `batch_size` | `1` | Training batch size |
203
+ | `learning_rate` | `5e-5` | Learning rate |
204
+ | `max_length` | `256` | Maximum sequence length |
205
+ | `output_dir` | `"./sft_finetuned_model"` | Output directory |
206
+ | `prompt_column` | `"prompt"` | Column name for prompts |
207
+ | `response_column` | `"response"` | Column name for responses |
208
+ | `messages_column` | `None` | Column with list of message dicts (multi-turn) |
209
+ | `train_on_completions_only` | `False` | Only compute loss on assistant tokens |
210
+ | `system_prompt` | `None` | System prompt for single-turn mode |
211
+ | `use_lora` | `False` | Enable LoRA |
212
+ | `lora_r` | `8` | LoRA rank |
213
+ | `lora_alpha` | `16` | LoRA alpha |
214
+ | `lora_dropout` | `0.1` | LoRA dropout |
215
+ | `quantize` | `True` | Enable 4-bit quantization |
216
+ | `quantization_type` | `"nf4"` | Quantization type (`nf4` or `fp4`) |
217
+ | `gradient_accumulation_steps` | `1` | Gradient accumulation steps |
218
+ | `warmup_steps` | `0` | LR scheduler warmup steps |
219
+ | `scheduler_type` | `"cosine"` | LR scheduler type |
220
+ | `max_grad_norm` | `1.0` | Gradient clipping norm (0 to disable) |
221
+ | `save_steps` | `0` | Save checkpoint every N steps (0 to disable) |
222
+ | `logging_steps` | `10` | Log metrics every N steps |
223
+
224
+ ### DPOConfig
225
+
226
+ Inherits all parameters from above, plus:
227
+
228
+ | Parameter | Default | Description |
229
+ |---|---|---|
230
+ | `beta` | `0.1` | DPO temperature parameter |
231
+ | `learning_rate` | `1e-6` | Learning rate (lower default for DPO) |
232
+ | `max_length` | `512` | Maximum sequence length |
233
+ | `chosen_column` | `"chosen"` | Column name for chosen responses |
234
+ | `rejected_column` | `"rejected"` | Column name for rejected responses |
235
+
236
+ ## Development
237
+
238
+ ```bash
239
+ git clone https://github.com/deduu/transfer.git
240
+ cd transfer
241
+ pip install -e ".[dev]"
242
+ pytest tests/ -v
243
+ ```
244
+
245
+ ## License
246
+
247
+ MIT
@@ -0,0 +1,217 @@
1
+ # Transfer
2
+
3
+ A modular PyTorch framework for fine-tuning Hugging Face language models.
4
+
5
+ Transfer simplifies applying techniques like Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO) to any Hugging Face causal language model.
6
+
7
+ ## Features
8
+
9
+ - **SFT and DPO** — two fine-tuning strategies out of the box
10
+ - **Multi-turn conversation support** — point to a `messages_column` of chat-style message lists
11
+ - **Completion-only loss masking** — train on assistant tokens only with `train_on_completions_only`
12
+ - **LoRA / QLoRA** — 4-bit quantization (nf4 / fp4) via bitsandbytes, configurable LoRA rank and target modules
13
+ - **Training loop controls** — gradient accumulation, LR scheduling (cosine, linear, constant, …), gradient clipping
14
+ - **Checkpointing and logging** — intermediate checkpoints every N steps, CSV training logs, optional W&B integration
15
+ - **Built-in evaluation** — perplexity and semantic entropy metrics
16
+ - **CLI** — `transfer train` and `transfer infer` commands for training and inference without writing Python
17
+
18
+ ## Installation
19
+
20
+ ### From PyPI
21
+
22
+ ```bash
23
+ pip install transfer-llm
24
+ ```
25
+
26
+ ### From source
27
+
28
+ ```bash
29
+ git clone https://github.com/deduu/transfer.git
30
+ cd transfer
31
+ pip install .
32
+ ```
33
+
34
+ ## Quick Start
35
+
36
+ ### Single-turn SFT
37
+
38
+ ```python
39
+ from datasets import Dataset
40
+ from transfer import Trainer, SFTConfig
41
+
42
+ # Prepare data
43
+ dataset = Dataset.from_dict({
44
+ "prompt": ["What is the capital of France?", "Explain gravity."],
45
+ "response": ["The capital of France is Paris.", "Gravity is ..."],
46
+ })
47
+
48
+ # Configure
49
+ config = SFTConfig(
50
+ model_name="google/gemma-2b",
51
+ num_epochs=3,
52
+ batch_size=2,
53
+ learning_rate=5e-5,
54
+ output_dir="./gemma-sft",
55
+ )
56
+
57
+ # Train
58
+ trainer = Trainer(task="sft", config=config, train_dataset=dataset)
59
+ trainer.train()
60
+ trainer.save_model()
61
+ ```
62
+
63
+ ### Multi-turn SFT with completion masking and LoRA
64
+
65
+ ```python
66
+ from datasets import Dataset
67
+ from transfer import Trainer, SFTConfig
68
+
69
+ # Each row contains a list of message dicts
70
+ conversations = [
71
+ {"messages": [
72
+ {"role": "system", "content": "You are a helpful assistant."},
73
+ {"role": "user", "content": "What is the weather in SF?"},
74
+ {"role": "assistant", "content": "It is 62 F and foggy."},
75
+ ]},
76
+ # ... more conversations
77
+ ]
78
+ dataset = Dataset.from_list(conversations)
79
+
80
+ config = SFTConfig(
81
+ model_name="meta-llama/Llama-3.2-3B-Instruct",
82
+ messages_column="messages", # read from chat-style column
83
+ train_on_completions_only=True, # loss on assistant tokens only
84
+ num_epochs=3,
85
+ batch_size=1,
86
+ learning_rate=2e-5,
87
+ max_length=1024,
88
+ gradient_accumulation_steps=4,
89
+ warmup_steps=10,
90
+ scheduler_type="cosine",
91
+ save_steps=50,
92
+ logging_steps=5,
93
+ use_lora=True,
94
+ lora_r=16,
95
+ lora_alpha=32,
96
+ lora_dropout=0.05,
97
+ output_dir="./agent-sft-output",
98
+ )
99
+
100
+ trainer = Trainer(task="sft", config=config, train_dataset=dataset)
101
+ trainer.train()
102
+ trainer.save_model()
103
+ ```
104
+
105
+ ### DPO
106
+
107
+ ```python
108
+ from datasets import load_dataset
109
+ from transfer import Trainer, DPOConfig
110
+
111
+ dataset = load_dataset("Anthropic/hh-rlhf", split="train[:1%]")
112
+
113
+ config = DPOConfig(
114
+ model_name="google/gemma-2b",
115
+ num_epochs=1,
116
+ batch_size=2,
117
+ beta=0.1,
118
+ output_dir="./gemma-dpo",
119
+ )
120
+
121
+ trainer = Trainer(task="dpo", config=config, train_dataset=dataset)
122
+ trainer.train()
123
+ trainer.save_model()
124
+ ```
125
+
126
+ ## CLI Usage
127
+
128
+ The `transfer` command is installed automatically with the package.
129
+
130
+ ### Training
131
+
132
+ ```bash
133
+ transfer train \
134
+ --task sft \
135
+ --model_name google/gemma-2b \
136
+ --dataset_path data.jsonl \
137
+ --output_dir ./my-sft-model \
138
+ --num_epochs 3 \
139
+ --batch_size 4 \
140
+ --learning_rate 2e-4 \
141
+ --use_lora \
142
+ --gradient_accumulation_steps 4 \
143
+ --scheduler_type cosine \
144
+ --warmup_steps 50 \
145
+ --save_steps 200 \
146
+ --logging_steps 10
147
+ ```
148
+
149
+ Use `--messages_column messages --train_on_completions_only` for multi-turn datasets.
150
+
151
+ ### Inference
152
+
153
+ ```bash
154
+ transfer infer \
155
+ --model_name google/gemma-2b \
156
+ --adapter_path ./my-sft-model \
157
+ --prompt "Explain quantum computing" \
158
+ --use_chat_template \
159
+ --max_new_tokens 256 \
160
+ --temperature 0.7 \
161
+ --do_sample
162
+ ```
163
+
164
+ ## Configuration
165
+
166
+ ### SFTConfig
167
+
168
+ | Parameter | Default | Description |
169
+ |---|---|---|
170
+ | `model_name` | `"google/gemma-2b"` | Hugging Face model name or path |
171
+ | `num_epochs` | `30` | Number of training epochs |
172
+ | `batch_size` | `1` | Training batch size |
173
+ | `learning_rate` | `5e-5` | Learning rate |
174
+ | `max_length` | `256` | Maximum sequence length |
175
+ | `output_dir` | `"./sft_finetuned_model"` | Output directory |
176
+ | `prompt_column` | `"prompt"` | Column name for prompts |
177
+ | `response_column` | `"response"` | Column name for responses |
178
+ | `messages_column` | `None` | Column with list of message dicts (multi-turn) |
179
+ | `train_on_completions_only` | `False` | Only compute loss on assistant tokens |
180
+ | `system_prompt` | `None` | System prompt for single-turn mode |
181
+ | `use_lora` | `False` | Enable LoRA |
182
+ | `lora_r` | `8` | LoRA rank |
183
+ | `lora_alpha` | `16` | LoRA alpha |
184
+ | `lora_dropout` | `0.1` | LoRA dropout |
185
+ | `quantize` | `True` | Enable 4-bit quantization |
186
+ | `quantization_type` | `"nf4"` | Quantization type (`nf4` or `fp4`) |
187
+ | `gradient_accumulation_steps` | `1` | Gradient accumulation steps |
188
+ | `warmup_steps` | `0` | LR scheduler warmup steps |
189
+ | `scheduler_type` | `"cosine"` | LR scheduler type |
190
+ | `max_grad_norm` | `1.0` | Gradient clipping norm (0 to disable) |
191
+ | `save_steps` | `0` | Save checkpoint every N steps (0 to disable) |
192
+ | `logging_steps` | `10` | Log metrics every N steps |
193
+
194
+ ### DPOConfig
195
+
196
+ Inherits all parameters from above, plus:
197
+
198
+ | Parameter | Default | Description |
199
+ |---|---|---|
200
+ | `beta` | `0.1` | DPO temperature parameter |
201
+ | `learning_rate` | `1e-6` | Learning rate (lower default for DPO) |
202
+ | `max_length` | `512` | Maximum sequence length |
203
+ | `chosen_column` | `"chosen"` | Column name for chosen responses |
204
+ | `rejected_column` | `"rejected"` | Column name for rejected responses |
205
+
206
+ ## Development
207
+
208
+ ```bash
209
+ git clone https://github.com/deduu/transfer.git
210
+ cd transfer
211
+ pip install -e ".[dev]"
212
+ pytest tests/ -v
213
+ ```
214
+
215
+ ## License
216
+
217
+ MIT
@@ -0,0 +1,58 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61.0"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "transfer-llm"
7
+ version = "0.1.0"
8
+ authors = [
9
+ { name="Dedy", email="dedy.ariansyah@gmail.com" },
10
+ ]
11
+ description = "A modular PyTorch framework for fine-tuning Hugging Face models."
12
+ readme = "README.md"
13
+ requires-python = ">=3.11"
14
+ classifiers = [
15
+ "Programming Language :: Python :: 3",
16
+ "License :: OSI Approved :: MIT License",
17
+ "Operating System :: OS Independent",
18
+ ]
19
+ dependencies = [
20
+ "numpy>=1.21.6",
21
+ "accelerate>=1.0.0",
22
+ "bitsandbytes>=0.44.0",
23
+ "transformers",
24
+ "datasets",
25
+ "tqdm",
26
+ "pyyaml",
27
+ "peft",
28
+ "sentence-transformers",
29
+ "scikit-learn",
30
+ "scipy",
31
+ "hf-xet>=1.2.0",
32
+ "torch>=2.3.0",
33
+ ]
34
+
35
+ [project.optional-dependencies]
36
+ dev = ["pytest>=7.0"]
37
+
38
+ [project.urls]
39
+ Homepage = "https://github.com/deduu/transfer"
40
+ Issues = "https://github.com/deduu/transfer/issues"
41
+
42
+ [project.scripts]
43
+ transfer = "transfer.cli:main"
44
+
45
+ [tool.pytest.ini_options]
46
+ testpaths = ["tests"]
47
+
48
+ [tool.setuptools]
49
+ packages = ["transfer", "transfer.strategies"]
50
+
51
+ [tool.uv.sources]
52
+ torch = { index = "pytorch-cu124" }
53
+
54
+ [[tool.uv.index]]
55
+ name = "pytorch-cu124"
56
+ url = "https://download.pytorch.org/whl/cu124"
57
+ explicit = true
58
+
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
File without changes
File without changes
File without changes
@@ -0,0 +1,141 @@
1
+ import torch
2
+ import gc
3
+ from math import exp
4
+ from tqdm import tqdm
5
+ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
6
+ from peft import PeftModel
7
+
8
+ PROMPT_TO_TEST = [
9
+ "Bagaimana cara menjelaskan Kepo kepada orang asing?",
10
+ # "Tulis artikel mengenai budaya kepo di Indonesia.",
11
+ # "Apakah kepo itu muncul secara genetik?"
12
+ ]
13
+
14
+ BASE_MODEL_NAME = "meta-llama/Llama-3.2-3B-Instruct"
15
+ LORA_ADAPTER_PATH = "./finetuned_model/llama-3.2-3b-it-kepo-lora"
16
+
17
+ bnb_config = BitsAndBytesConfig(
18
+ load_in_4bit=True,
19
+ bnb_4bit_quant_type="nf4",
20
+ bnb_4bit_compute_dtype=torch.bfloat16,
21
+ bnb_4bit_use_double_quant=True,
22
+ )
23
+
24
+
25
+ def compute_per_prompt(model, tokenizer, prompt, device, log_details=False):
26
+ enc = tokenizer(prompt, return_tensors="pt", truncation=True).to(device)
27
+
28
+ with torch.inference_mode():
29
+ outputs = model(**enc, labels=enc["input_ids"])
30
+ loss = outputs.loss.item()
31
+
32
+ if log_details:
33
+ logits = outputs.logits[:, :-1, :]
34
+ labels = enc["input_ids"][:, 1:]
35
+ log_probs = torch.log_softmax(logits, dim=-1)
36
+ token_log_probs = log_probs.gather(-1,
37
+ labels.unsqueeze(-1)).squeeze(-1)
38
+
39
+ # Convert each token's loss
40
+ token_losses = (-token_log_probs.squeeze()).tolist()
41
+
42
+ print(f"\n🧠 Token losses for: `{prompt}`")
43
+ tokens = tokenizer.convert_ids_to_tokens(enc["input_ids"].squeeze())
44
+ for t, l in zip(tokens[1:], token_losses): # skip first token
45
+ print(f"{t:20s} loss={l:.3f}")
46
+
47
+ return exp(loss), loss
48
+
49
+
50
+ def compute_token_weighted(model, tokenizer, device, log_details=False):
51
+ """Compute total token-weighted perplexity."""
52
+ total_loss = 0.0
53
+ total_tokens = 0
54
+
55
+ for prompt in tqdm(PROMPT_TO_TEST, desc="Token-weighted scoring"):
56
+ enc = tokenizer(prompt, return_tensors="pt",
57
+ truncation=True).to(device)
58
+ with torch.inference_mode():
59
+ outputs = model(**enc, labels=enc["input_ids"])
60
+
61
+ logits = outputs.logits[:, :-1, :]
62
+ labels = enc["input_ids"][:, 1:]
63
+ log_probs = torch.log_softmax(logits, dim=-1)
64
+ token_log_probs = log_probs.gather(-1,
65
+ labels.unsqueeze(-1)).squeeze(-1)
66
+
67
+ token_losses = -token_log_probs
68
+ total_loss += token_losses.sum().item()
69
+ total_tokens += token_losses.numel()
70
+
71
+ if log_details:
72
+ print(f"\n📍 Prompt: {prompt}")
73
+ print(
74
+ f" Tokens: {token_losses.numel()} | Sum Loss: {token_losses.sum().item():.3f}")
75
+
76
+ avg_loss = total_loss / total_tokens
77
+ return exp(avg_loss), avg_loss
78
+
79
+
80
+ def evaluate(label, tokenizer, load_lora=False, method="prompt_avg", log_details=False):
81
+ device = "cuda" if torch.cuda.is_available() else "cpu"
82
+
83
+ print(f"\n🔄 Loading {label}...")
84
+ model = AutoModelForCausalLM.from_pretrained(
85
+ BASE_MODEL_NAME,
86
+ quantization_config=bnb_config,
87
+ device_map="auto"
88
+ )
89
+
90
+ if load_lora:
91
+ print("🔌 Applying LoRA adapter...")
92
+ model = PeftModel.from_pretrained(model, LORA_ADAPTER_PATH)
93
+
94
+ model.eval()
95
+
96
+ if method == "prompt_avg":
97
+ ppl_list = []
98
+ for prompt in PROMPT_TO_TEST:
99
+ ppl, _ = compute_per_prompt(
100
+ model, tokenizer, prompt, device, log_details)
101
+ ppl_list.append(ppl)
102
+ score = sum(ppl_list) / len(ppl_list)
103
+
104
+ elif method == "token_weighted":
105
+ score, _ = compute_token_weighted(
106
+ model, tokenizer, device, log_details)
107
+
108
+ print(f"\n📍 Perplexity ({method}) for {label}: {score:.4f}")
109
+
110
+ del model
111
+ torch.cuda.empty_cache()
112
+ gc.collect()
113
+ print("🧹 GPU cache cleared")
114
+
115
+ return score
116
+
117
+
118
+ def main():
119
+ tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
120
+ tokenizer.pad_token = tokenizer.eos_token
121
+
122
+ print("\n📦 4-bit quantized evaluation starting...\n")
123
+
124
+ # Toggle here:
125
+ method = "prompt_avg" # or: "token_weighted"
126
+ log_details = True # show token losses
127
+
128
+ base = evaluate("Base Model", tokenizer, load_lora=False,
129
+ method=method, log_details=log_details)
130
+ lora = evaluate("LoRA Model", tokenizer, load_lora=True,
131
+ method=method, log_details=log_details)
132
+
133
+ print("\n============== FINAL RESULTS ==============")
134
+ print(f"Base model PPL: {base:.4f}")
135
+ print(f"LoRA model PPL: {lora:.4f}")
136
+ print(f"Improvement: {((base - lora)/base) * 100:.2f}%")
137
+ print("===========================================")
138
+
139
+
140
+ if __name__ == "__main__":
141
+ main()