Create ft.py

ft.py

@@ -0,0 +1,116 @@
+#fine-tuning code
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainerCallback
+from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model
+from datasets import Dataset
+import json
+from trl import SFTTrainer, SFTConfig
+from transformers import TrainingArguments
+import time
+import os
+
+class CustomCallback(TrainerCallback):
+    def __init__(self):
+        self.start_time = time.time()
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        print("Training has begun!")
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if state.global_step % args.logging_steps == 0:
+            elapsed_time = time.time() - self.start_time
+            if state.log_history:
+                loss = state.log_history[-1].get('loss', 0)
+                print(f"Step: {state.global_step}, Loss: {loss:.4f}, Time: {elapsed_time:.2f}s")
+            else:
+                print(f"Step: {state.global_step}, Loss: N/A, Time: {elapsed_time:.2f}s")
+
+    def on_train_end(self, args, state, control, **kwargs):
+        print("Training has ended!")
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "7"  
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+model_path = "Google/gemma-2-9b-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+tokenizer.pad_token = tokenizer.eos_token
+
+# (QLORA)
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    quantization_config=bnb_config,
+    device_map="auto",  
+    torch_dtype=torch.float16,
+)
+
+# PEFT 
+model = prepare_model_for_kbit_training(model)
+lora_config = LoraConfig(
+    r=6,
+    lora_alpha=8,
+    lora_dropout=0.05,
+    target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
+    task_type="CAUSAL_LM",
+)
+model = get_peft_model(model, lora_config)
+
+# data
+with open('en_ko_data', 'r', encoding='utf-8') as f:
+    data = json.load(f)
+
+# prompt
+def generate_prompt(en_text, ko_text):
+    return f"""<bos><start_of_turn>user
+Please translate the following English colloquial expression into Korean.:
+
+{en_text}<end_of_turn>
+<start_of_turn>model
+{ko_text}<end_of_turn><eos>"""
+
+
+key = list(data.keys())[0]  
+dataset = [{"text": generate_prompt(item['en_original'], item['ko'])} for item in data[key]]
+dataset = Dataset.from_list(dataset)
+
+# training set
+training_args = TrainingArguments(
+    output_dir="./results",
+    num_train_epochs=3,
+    per_device_train_batch_size=4,
+    gradient_accumulation_steps=4,
+    save_steps=100,
+    logging_steps=1,  
+    learning_rate=2e-4,
+    weight_decay=0.01,
+    fp16=True,
+    optim="paged_adamw_8bit",
+)
+
+# SFTTrainer 
+trainer = SFTTrainer(
+    model=model,
+    train_dataset=dataset,
+    args=training_args,
+    tokenizer=tokenizer,
+    dataset_text_field="text",
+    max_seq_length=512,
+)
+
+
+trainer.add_callback(CustomCallback())
+
+
+# train
+trainer.train()
+
+# save
+trainer.save_model("./gemma2_9b_ko_translator")