Ssul's Blog

AI Product 개발전략과 개발기 - text-classification모델 파인튜닝 본문

AI & ML/사용하기

AI Product 개발전략과 개발기 - text-classification모델 파인튜닝

Ssul 2024. 6. 27. 23:11

AI스팸분류기를 제작하고 있다.

지난번에는 text-generation모델을 기반으로 instruction-finetuning을 통한 모델을 개발했다.

이번 글에서는 text-classification모델을 기반으로 파인튜닝하여 모델을 개발해볼 예정이다.

(*text-classification의 대표적인 예는 문장의 긍정/부정을 예측하는 모델이 있다.)

 

1-2는 기존글과 동일하기에 생략

https://issul.tistory.com/455

 

3. text-classfication모델 가져와서, 양자화하고, LoRA붙이기

device = 'auto' #@param {type: "string"}
base_LLM_model = 'yanolja/EEVE-Korean-Instruct-2.8B-v1.0' #@param {type: "string"}

# 토크나이저 및 모델 로드
model = AutoModelForSequenceClassification.from_pretrained(
    base_LLM_model,
    # load_in_8bit=True, # LoRA
    # torch_dtype=torch.float16,
    load_in_4bit=True, # Quantization Load
    num_labels=6,
    device_map=device)

tokenizer = AutoTokenizer.from_pretrained(base_LLM_model)

- 기존과 다른점은 'AutoModelForSequenceClassification으로 가져오는 것이다. 옵션으로 num_labels=6을 준다. 이유는 내가 분류할 문자의 종류가 6개이기 때문이다. 만약에 3종류로 분류한다 하면, num_labels=3으로 하면 된다.

 

###########################################
# 파인튜닝 고고
# NF4 양자화를 위한 설정
lora_config = LoraConfig(
    r=4, # LoRA 가중치 행렬의 rank. 정수형이며 값이 작을수록 trainable parameter가 적어짐
    lora_alpha=8, # LoRA 스케일링 팩터. 추론 시 PLM weight와 합칠 때 LoRA weight의 스케일을 일정하게 유지하기 위해 사용
    lora_dropout=0.05,
    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'], # LoRA를 적용할 layer. 모델 아키텍처에 따라 달라짐
    bias='none', # bias 파라미터를 학습시킬지 지정. ['none', 'all', 'lora_only']
    task_type=TaskType.CAUSAL_LM
)

# 양자화된 모델을 학습하기 전, 전처리를 위해 호출
model = prepare_model_for_kbit_training(model)
# LoRA 학습을 위해서는 아래와 같이 peft를 사용하여 모델을 wrapping 해주어야 함
model = get_peft_model(model, lora_config)

# 학습 파라미터 확인
model.print_trainable_parameters()

print(model)

- LoRA옵션을 설정하고, LoRA를 붙인다.

- 실제 모델을 출력해보면, LoRA가 붙어있는것과, 제일마지막이 6개로 됨을 확인할 수 있다.

 

4. 데이터셋 구성하기

- text-classification이기 때문에, text로 구성된 라벨을 모두 숫자로 변경

import pandas as pd

# CSV 파일 경로를 지정합니다.
file_path = '/content/drive/MyDrive/00_AI연구/02_스팸분류기Proj/spam_dataset_final2.csv'

# CSV 파일을 pandas DataFrame으로 불러옵니다.
dataset = pd.read_csv(file_path)
# dataset = pd.read_csv(file_path, encoding='utf-8')

# DataFrame 출력
print(dataset)
# 라벨 매핑
label_mapping = {
    '스팸문자': 0,
    '프로모션문자': 1,
    '정부문자': 2,
    '정치선거문자': 3,  # '정치문자'와 '정치선거문자'를 동일하게 처리
    '개인문자': 4,
    '인증문자': 5
}

dataset['num_label'] = dataset['label'].map(label_mapping)

- 기존의 텍스트 라벨을 숫자 라벨로 변환

 

import re
# 텍스트 정리 함수 정의
def clean_text(text):
    text = re.sub(r'\[.*?\]', '', text)  # 대괄호 안의 텍스트 제거
    text = re.sub(r'\(.*?\)', '', text)  # 소괄호 안의 텍스트 제거
    text = re.sub(r'\n', ' ', text)  # 줄바꿈 문자 제거
    text = re.sub(r'\s+', ' ', text)  # 여러 개의 공백을 하나의 공백으로 대체
    text = re.sub(r'[^가-힣a-zA-Z\s]', '', text)  # 한글, 영어 알파벳, 공백을 제외한 문자 제거
    return text.strip()

# 'sms' 컬럼에 정리 함수 적용
dataset['cleaned_sms'] = dataset['sms'].apply(clean_text)

- sms문자열 Preprocessing

 

import pandas as pd
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from sklearn.model_selection import train_test_split

train_df, test_df = train_test_split(dataset, test_size=0.2)
# Hugging Face 데이터셋으로 변환
train_dataset = Dataset.from_pandas(train_df)
test_dataset = Dataset.from_pandas(test_df)
def tokenize_function(examples):
    return tokenizer(examples['cleaned_sms'], padding='max_length', truncation=True)

tokenized_train_dataset = train_dataset.map(tokenize_function, batched=True)
tokenized_test_dataset = test_dataset.map(tokenize_function, batched=True)

# 'label'을 'text_label'로 변경
tokenized_train_dataset = tokenized_train_dataset.rename_column("label", "text_label")
tokenized_test_dataset = tokenized_test_dataset.rename_column("label", "text_label")
tokenized_train_dataset

# 'num_label'을 'label'로 변경
tokenized_train_dataset = tokenized_train_dataset.rename_column("num_label", "labels")
tokenized_test_dataset = tokenized_test_dataset.rename_column("num_label", "labels")
tokenized_train_dataset

Dataset({
    features: ['sms', 'text_label', 'labels', 'cleaned_sms', '__index_level_0__', 'input_ids', 'attention_mask'],
    num_rows: 800
})

- 실제 모델 학습시 사용되는 데이터는 Input_ids, attention_mask, labels

- 입력으로 Input_ids, attention_mask이 되고, 출력된 예측값과 labels를 비교해서 업데이트 하는 구조

 

 

5. 학습 실행

# 트레이너 설정
training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    num_train_epochs=1,
    weight_decay=0.01,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train_dataset,
    eval_dataset=tokenized_test_dataset,
    tokenizer=tokenizer,
)

# 모델 훈련
trainer.train()