File size: 2,615 Bytes
6a08381 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
# This script exemplify a selective fine-tuning method based on the condition number using freely available data and LLM
# Author: Oswaldo Ludwig (now with AI support)
# Date: 03/07/2025
# In case of publication using this script or ideas in this script, cite:
# Ludwig, Oswaldo. "The Condition Number as a Scale-Invariant Proxy for Information Encoding in Neural Units." arXiv preprint arXiv:2506.16289 (2025).
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from selective_fine_tuning import SelectiveFineTuningOptimizer
# Dataset using AG News
class RealTextDataset(Dataset):
def __init__(self, tokenizer, split='train', max_samples=200, seq_len=64):
dataset = load_dataset("ag_news", split=split)
self.samples = dataset.select(range(max_samples))
self.tokenizer = tokenizer
self.seq_len = seq_len
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
text = self.samples[idx]['text']
encoding = self.tokenizer(text, truncation=True, padding='max_length', max_length=self.seq_len, return_tensors='pt')
input_ids = encoding['input_ids'].squeeze(0)
return input_ids, input_ids.clone()
# Training loop
def train():
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
print("Tokenizer and LLM loaded", flush=True)
dataset = RealTextDataset(tokenizer=tokenizer, max_samples=200, seq_len=64)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
print("Data loader and dataset loaded", flush=True)
criterion = nn.CrossEntropyLoss()
optimizer_wrapper = SelectiveFineTuningOptimizer(
model=model,
base_optimizer_cls=optim.AdamW,
optimizer_args={'lr': 5e-5},
condition_file='condition_numbers.json',
num_tensors_to_finetune=20,
recompute=True
)
print("Optimizer instantiated", flush=True)
model.train()
for epoch in range(3):
total_loss = 0
for inputs, targets in dataloader:
optimizer_wrapper.zero_grad()
outputs = model(inputs, labels=targets)
loss = outputs.loss
loss.backward()
optimizer_wrapper.step()
total_loss += loss.item()
print(f"Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}")
if __name__ == '__main__':
train()
|