Maaac/CodeLLaMA-Linux-BugFix

1# QLoRA fine-tuning for CodeLLaMA-7B-Instruct
2# Requirements: transformers, peft, accelerate, bitsandbytes, datasets
3
4from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, BitsAndBytesConfig
5from peft import LoraConfig, get_peft_model
6from datasets import load_dataset
7import torch
8import os
9
10# Paths and parameters
11BASE_MODEL = "codellama/CodeLlama-7b-Instruct-hf"
12DATA_PATH = "../dataset/training_data_prompt_completion.jsonl"
13OUTPUT_DIR = "./output/qlora-codellama-bugfix"
14
15# Load dataset (prompt, completion)
16dataset = load_dataset("json", data_files=DATA_PATH, split="train")
17
18# Apply formatting for supervised fine-tuning
19def format(example):
20    prompt = tokenizer(
21        example["prompt"],
22        truncation=True,
23        padding="max_length",
24        max_length=512
25    )
26    completion = tokenizer(
27        example["completion"],
28        truncation=True,
29        padding="max_length",
30        max_length=512
31    )
32    input_ids = prompt["input_ids"] + completion["input_ids"]
33    labels = [-100] * len(prompt["input_ids"]) + completion["input_ids"]
34
35    return {
36        "input_ids": input_ids[:1024],
37        "labels": labels[:1024]
38    }
39
40# Load tokenizer and base model
41bnb_config = BitsAndBytesConfig(
42    load_in_4bit=True,
43    bnb_4bit_use_double_quant=True,
44    bnb_4bit_quant_type="nf4",
45    bnb_4bit_compute_dtype=torch.float16
46)
47
48tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
49tokenizer.pad_token = tokenizer.eos_token  # Required for padding
50
51model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, quantization_config=bnb_config, device_map="auto")
52
53# Apply QLoRA
54lora_config = LoraConfig(
55    r=64,
56    lora_alpha=16,
57    lora_dropout=0.1,
58    bias="none",
59    task_type="CAUSAL_LM"
60)
61
62model = get_peft_model(model, lora_config)
63
64# Tokenize dataset
65dataset = dataset.map(format, remove_columns=["prompt", "completion"])
66
67# Training args
68training_args = TrainingArguments(
69    output_dir=OUTPUT_DIR,
70    num_train_epochs=3,
71    per_device_train_batch_size=1,
72    gradient_accumulation_steps=4,
73    learning_rate=2e-4,
74    logging_dir=f"{OUTPUT_DIR}/logs",
75    logging_steps=10,
76    save_strategy="epoch",
77    bf16=False,
78    fp16=True,
79    save_total_limit=2,
80    report_to="none",
81    push_to_hub=False
82)
83
84# Trainer
85trainer = Trainer(
86    model=model,
87    args=training_args,
88    train_dataset=dataset,
89    tokenizer=tokenizer
90)
91
92trainer.train()
93
94model.save_pretrained(OUTPUT_DIR)
95tokenizer.save_pretrained(OUTPUT_DIR)
96print(f"[DONE] Model saved to {OUTPUT_DIR}")
97