Upload Training Code

PEFT_LoRA_CLM_KULLM3_TrainingCode.py

@@ -0,0 +1,182 @@
+import json
+import os
+from tqdm import tqdm
+import pandas as pd
+import numpy as np
+import re
+import transformers
+from peft import LoraConfig, PeftConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
+from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, TrainingArguments, Trainer, default_data_collator
+import torch
+from dataclasses import dataclass, field
+from typing import Optional
+from dataclass_csv import DataclassReader
+from torch.utils.data import DataLoader
+from datasets import Dataset
+from trl import DataCollatorForCompletionOnlyLM
+from peft import get_peft_model, LoraConfig, TaskType
+import sys
+
+import time
+from enum import Enum
+
+# ignore warning
+import warnings
+warnings.filterwarnings("ignore")
+
+os.environ['HF_TOKEN'] = "SET YOUR HUGGINGFACE ACCESS TOKEN" # it should be write access-token
+os.environ["HUGGINGFACEHUB_API_TOKEN"] = "SET YOUR HUGGINGFACE ACCESS TOKEN" # write access-token
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+
+class SpecialTokens(str, Enum):
+    soc_token = "<|startofcall|>" # 첫 대화의 시작이 상담원이면 안 됨 -> 첫 대화의 시작이 상담원인 경우 human 메세지로 soc 토큰 삽입
+    # interjection_token = "<|interjection|>" # 화자가 연달아 2번 발화할 때 중간에 추임새 토큰을 넣어준다.
+    privacy_token = "<|private|>" # 원본에서 xx, ㅇㅇㅇ, 00 등 개인정보로 표시된 부분을 프라이버시 토큰으로 변환
+    pad_token = "[PAD]"
+
+    @classmethod
+    def list(cls):
+        return [c.value for c in cls]
+
+class KuLLM3:
+    def __init__(self, device, model_path="nlpai-lab/KULLM3"):
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map=device
+        )
+        # Special token들을 tokenizer에 추가
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            pad_token=SpecialTokens.pad_token.value,
+            padding_side='right',
+            additional_special_tokens=SpecialTokens.list(),
+            model_max_length=30000
+        )
+        self.tokenizer.chat_template = '{% if messages[0][\'role\'] == \'system\' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0][\'content\'] %}{% else %}{% set loop_messages = messages %}{% set system_message = "당신은 고객의 전화에 응대하는 전화 상담원입니다.\n 당신은 비도덕적이거나, 성적이거나, 불법적이거나 또는 사회 통념적으로 허용되지 않는 발언은 하지 않습니다.\n 고객에게 친절하게 대화하며, 고객의 응답에 가능한 정확하고 예의 바르게 응답함으로써 최대한 도와주려고 노력합니다.\n 고객의 질문을 이해하지 못했다면, 어떤 부분을 이해하지 못했는지 설명하고 고객에게 구체적인 질문을 요구합니다.\n 당신은 고객과 전화로 소통하기 때문에 답변이 간결해야 합니다. 거짓 정보를 발언하지 않도록 주의합니다." %}{% endif %}{% for message in loop_messages %}{% if (message[\'role\'] == \'user\') != (loop.index0 % 2 == 0) %}{{ raise_exception(\'Conversation roles must alternate user/assistant/user/assistant/...\') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = \'<<SYS>>\\n\' + system_message + \'\\n<</SYS>>\\n\\n\' + message[\'content\'] %}{% else %}{% set content = message[\'content\'] %}{% endif %}{% if message[\'role\'] == \'user\' %}{{ bos_token + \'[INST] \' + content.strip() + \' [/INST]\'}}{% elif message[\'role\'] == \'system\' %}{{ \'<<SYS>>\\n\' + content.strip() + \'\\n<</SYS>>\\n\\n\' }}{% elif message[\'role\'] == \'assistant\' %}{{ \' \'  + content.strip() + \' \' + eos_token }}{% endif %}{% endfor %}'
+        # self.tokenizer = AutoTokenizer.from_pretrained(
+        #     "song9/KuLLM3-LoRA-CC"
+        # )
+        # embedding layer의 input을 기존 dimension에서 special token개수만큼 추가한 차원만큼 resize
+        self.model.resize_token_embeddings(len(self.tokenizer))
+
+        # LoRA Linear 레이어를 생성한다.
+        self.config = LoraConfig(
+            task_type=TaskType.CAUSAL_LM,
+            inference_mode=False,
+            r=64, 
+            lora_alpha=128, 
+            lora_dropout=0.01, 
+            target_modules=["embed_tokens", "lm_head", "q_proj", "v_proj"]
+        )
+        # model에 LoRA 레이어를 추가한다.
+        self.model = get_peft_model(self.model, self.config)
+
+if __name__ == '__main__':
+    ### 금요일 오후 6시반부터 시작
+    print("✅Start Timer", flush=True)
+    time.sleep(381600)
+    
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    kullm3 = KuLLM3(device=device)
+    tokenizer, model = kullm3.tokenizer, kullm3.model
+    
+    root_dir = os.path.join('SET PATH TO your dataset dir path')
+    
+    # save as tsv file
+    # Concat two dataset(AI-Hub/민원 and AI-Hub/용도별), make train_combined.tsv and valid_combined.tsv
+    training = os.path.join(root_dir, 'train_combined.tsv')
+    validation = os.path.join(root_dir, 'valid_combined.tsv')
+    # df.to_csv(savepath, sep='\t', index=False)
+    # load as tsv file
+    train_df = pd.read_csv(training, sep='\t')
+    val_df = pd.read_csv(validation, sep='\t')
+    
+    response_template = '[/INST]'
+    instruction_template = '[INST]'
+    print('✅Preparing dataset', flush=True)
+    train_dataset = Dataset.from_list(train_df['dialogue'].apply(lambda x:tokenizer(
+        x, return_length=False, return_attention_mask=True
+    )).to_list())
+    val_dataset = Dataset.from_list(val_df['dialogue'].apply(lambda x:tokenizer(
+        x, return_length=False, return_attention_mask=True
+    )).to_list())
+    collator = DataCollatorForCompletionOnlyLM(
+        instruction_template=instruction_template,
+        response_template=response_template,
+        tokenizer=tokenizer
+    )
+
+    ### 금요일 오후 6시반부터 시작
+    # print("✅Start Timer", flush=True)
+    # time.sleep(381600) 
+    
+    peft_model_path = "SET YOUR HUGGINGFACE models path" # ex) song9/CC-KuLLM3-LoRA
+    print('✅Start Training', flush=True)
+    # train args
+    batchsize_per_step = 2
+    save_steps = 8000
+    output_dirpath = 'SET YOUR CHECKPOINT DIR PATH'
+    training_args = TrainingArguments(
+        output_dir=output_dirpath, # 모델과 체크포인트를 저장할 디렉토리
+        save_total_limit=1, # 저장할 체크포인트의 최대 개수
+        load_best_model_at_end=True, # 훈련 종료 시 가장 좋은 성능을 낸 모델을 로드함
+        save_steps=save_steps, # load_best_model_at_end가 True일 때 설정해야 한다. eval_steps, logging_steps와 값을 맞춰야 한다.
+        logging_steps=int(save_steps/10), # 학습 중 로그를 출력할 step 수.
+        # Train 관련 파라미터
+        num_train_epochs=2, # 훈련 epoch 수
+        per_device_train_batch_size=batchsize_per_step, # 각 device(gpu) 당 배치 사이즈. 
+        remove_unused_columns=False, # 데이터셋에서 사용하지 않는 열이 있어도 제거하지 않는다.
+        bf16=True, # dtype이 bfloat16 형식인지 여부
+        dataloader_drop_last=True, # ex) True: batch_size가 2인데 마지막 batch엔 데이터가 1개뿐이면 해당 batch는 버린다.
+        group_by_length=True, # 이걸 해야 길이가 비슷한 데이터끼리 같은 batch에 묶임.
+        # 최적화 관련    
+        gradient_checkpointing=True, # 메모리 절약을 위해 역전파 시 체크포인트를 사용하여 일부 중간 계산을 재계산하는 방식이다.
+        gradient_checkpointing_kwargs={"use_reentrant": False},  # Gradient Checkpointing 동작 방식을 설정하는 옵션 (PyTorch의 비권장 reentrant API 사용을 피함)
+        # gradient_accumulation_steps=16, # batch_size와 반비례하는 파라미터. batch size를 키울 수 없을 때, 여러 개의 batch의 gradient를 축적하여 마치 하나의 큰 batch를 사용하는 것 같은 효과를 냄. 그러나 메모리에 부담이 됨.
+        torch_empty_cache_steps=batchsize_per_step, # 설정 스텝마다 `torch.cuda.empty_cache()`를 호출하여 메모리 캐시를 비움 (메모리 누수 방지)
+        # dataloader_num_workers=1, # data를 load할 때 사용할 process의 수. default=0 이면 main process에서 dataloader를 수행 (사용하면 메모리 사용량 급증...)
+        # torch_compile=True, # torch.compile을 사용해 속도 향상
+        # Optimizer 및 Scheduler
+        lr_scheduler_type='linear',
+        # warmup_steps=200, # 학습 초기 단계에서 설정값 단계까지 학습률을 서서히 증가시킨다.
+        warmup_ratio=0.05, # 전체 step을 모를 때 warmup을 전체 step 중 ratio만큼 적용
+        # optim="adamw_torch", # AdamW 옴티마이저 사용
+        # optim_args={ ## ? 이게 아닌가봄 에러 유발.
+        #     'lr':5e-5,'weight_decay':0.01
+        # },    
+        weight_decay=0.01, # L2 규제용 가중치 감소율이다. 과적합 방지를 위해 사용하는 정규화 기술.
+        learning_rate=1e-5, # 학습률
+        
+        # # eval
+        per_device_eval_batch_size=batchsize_per_step,
+        eval_strategy='steps', 
+        # 평가를 실행할 전략 ('steps'로 설정 시 eval_steps마다 평가가 수행됨)
+        # 'epoch'로 설정 시 train epoch 종료 시마다 eval을 수행한다.
+        eval_steps=save_steps, # 평가를 수행할 스텝 간격
+        
+        # # huggingface에 모델을 업로드할 때 필요한 매개변수
+        hub_model_id=peft_model_path,
+        push_to_hub=True,
+        hub_private_repo=False  
+    )
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+        data_collator=collator
+    )
+    
+    trainer.train()
+
+    output_dirpath_backup = output_dirpath+'_backup'
+    if os.path.exists(output_dirpath_backup):
+        os.mkdir(output_dirpath_backup)
+    model.save_pretrained(output_dirpath_backup)
+    tokenizer.save_pretrained(output_dirpath_backup)
+    # trainer.push_to_hub()
+    # trainer.model.push_to_hub(training_args.output_dir)
+
+    # 직접 huggingface-cli로 push