lora-fine-tuning/lora_example.py at main · bhimanbaghel/lora-fine-tuning · GitHub

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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import os
import gc
from src.args import parse_args, read_hparam
from src.models import get_model_and_tokenizer, get_peft_config, get_fine_tuned_model
from src.dataset_preparation import prepare_dataset
from src.training import custom_train
from src.evaluation import evaluate
from peft import get_peft_model

def main():
    args = parse_args()
    args.hparam = read_hparam(args.hparam_path)
    print("Hparam Loaded")

    # Create output directories
    os.makedirs(args.save_dir, exist_ok=True)

    # Get model and tokenizer
    model, tokenizer = get_model_and_tokenizer(args)
    print("Model Loaded")

    # Prepare dataset
    dataset, data_collator = prepare_dataset(tokenizer, args)
    print("Dataset Loaded")

    # Print samples from the training set
    print("===== TRAINING SAMPLES =====")
    for i in range(2):  # Print first 2 samples
        sample = dataset['train'][i]
        print(f"Sample {i+1}:")

        # Decode input_ids back to text
        input_text = tokenizer.decode(sample['input_ids'], skip_special_tokens=True)
        print(f"Input: {input_text}")

        # Decode labels back to text, handling -100 values
        label_ids = [id if id != -100 else tokenizer.pad_token_id for id in sample['labels']]
        label_text = tokenizer.decode(label_ids, skip_special_tokens=True)
        print(f"Label: {label_text}")
        print("-" * 50)

    # Print samples from the validation set
    print("\n===== VALIDATION SAMPLES =====")
    for i in range(2):  # Print first 2 samples
        sample = dataset['validation'][i]
        print(f"Sample {i+1}:")
        input_text = tokenizer.decode(sample['input_ids'], skip_special_tokens=True)
        print(f"Input: {input_text}")
        label_ids = [id if id != -100 else tokenizer.pad_token_id for id in sample['labels']]
        label_text = tokenizer.decode(label_ids, skip_special_tokens=True)
        print(f"Label: {label_text}")
        print("-" * 50)

    # Print samples from the test set
    print("\n===== TEST SAMPLES =====")
    for i in range(2):  # Print first 2 samples
        sample = dataset['test'][i]
        print(f"Sample {i+1}:")
        input_text = tokenizer.decode(sample['input_ids'], skip_special_tokens=True)
        print(f"Input: {input_text}")
        label_ids = [id if id != -100 else tokenizer.pad_token_id for id in sample['labels']]
        label_text = tokenizer.decode(label_ids, skip_special_tokens=True)
        print(f"Label: {label_text}")
        print("-" * 50)

    # Configure LoRA
    peft_config = get_peft_config(args)
    model = get_peft_model(model, peft_config)

    # Print trainable parameters
    model.print_trainable_parameters()

    # Train the model
    if "train" in dataset and not args.eval:
        custom_train(model, tokenizer, dataset, data_collator, args)
    else:
        print("Eval Mode, skipping training.")

    if args.eval:
        model = get_fine_tuned_model(args, model)
        evaluate(model, tokenizer, dataset, data_collator, args)


if __name__ == "__main__":
    main()