[28 일차] ABC 부트캠프 : 딥러닝 기초(2)

import os
import json
import glob
from omegaconf import OmegaConf

# 함수 정의: 최신 모델을 가져오는 함수
def get_latest_model():  
    nemo_model_paths = glob.glob('nemo_experiments/TextClassification/*/checkpoints/*.nemo')
    # 최신 순으로 정렬
    nemo_model_paths.sort(reverse=True)
    return nemo_model_paths[0]

# 데이터 디렉토리 설정
DATA_DIR = '/dli/task/data/federalist_papers_HM'

# 데이터 디렉토리 내 파일 목록 출력
os.system(f'ls {DATA_DIR}')

# train.tsv 및 dev.tsv 파일을 nemo 형식으로 수정하고 저장하는 부분
# 이 부분에서 기존 train.tsv와 dev.tsv 파일을 읽고 필요에 따라 수정한 뒤,
# train_nemo_format.tsv와 dev_nemo_format.tsv로 저장해야 합니다.
# 예를 들어, 아래는 간단한 형식 수정 예시입니다. (데이터 형식에 따라 다름)
train_data = []
dev_data = []

with open(f'{DATA_DIR}/train.tsv') as f:
    for line in f:
        label, text = line.strip().split('\t')
        train_data.append(f'{label}\t{text}')

with open(f'{DATA_DIR}/train_nemo_format.tsv', 'w') as f:
    for line in train_data:
        f.write(line + '\n')

with open(f'{DATA_DIR}/dev.tsv') as f:
    for line in f:
        label, text = line.strip().split('\t')
        dev_data.append(f'{label}\t{text}')

with open(f'{DATA_DIR}/dev_nemo_format.tsv', 'w') as f:
    for line in dev_data:
        f.write(line + '\n')

# 수정된 train 및 dev 파일의 샘플 출력
print("*****\ntrain_nemo_format.tsv sample\n*****")
os.system(f'head -n 3 {DATA_DIR}/train_nemo_format.tsv')
print("\n\n*****\ndev_nemo_format.tsv sample\n*****")
os.system(f'head -n 3 {DATA_DIR}/dev_nemo_format.tsv')

# train_nemo_format.tsv와 dev_nemo_format.tsv 파일에서 첫 두 줄을 읽어 평가를 위해 저장
step1 = []
try:
    with open(os.path.join(DATA_DIR,'train_nemo_format.tsv')) as f:
        content = f.readlines()
        step1 += content[:2]
    with open(os.path.join(DATA_DIR,'dev_nemo_format.tsv')) as f:
        content = f.readlines()
        step1 += content[:2]
except:
    pass
                
with open("my_assessment/step1.json", "w") as outfile: 
    json.dump(step1, outfile) 

# 기본 모델 설정 파일의 내용을 불러오기
CONFIG_DIR = "/dli/task/nemo/examples/nlp/text_classification/conf"
CONFIG_FILE = "text_classification_config.yaml"

config = OmegaConf.load(CONFIG_DIR + "/" + CONFIG_FILE)
print(OmegaConf.to_yaml(config.model))

# 사용 가능한 BERT 유사 언어 모델 목록 가져오기
from nemo.collections import nlp as nemo_nlp
print(nemo_nlp.modules.get_pretrained_lm_models_list())

# 설정 값 정의
NUM_CLASSES = 2  # 클래스 수, HAMILTON과 MADISON 두 개의 클래스가 있음
MAX_SEQ_LENGTH = 256  # 최대 시퀀스 길이
BATCH_SIZE = 16  # 배치 크기

# 데이터 파일 경로 설정
PATH_TO_TRAIN_FILE = "/dli/task/data/federalist_papers_HM/train_nemo_format.tsv"
PATH_TO_VAL_FILE = "/dli/task/data/federalist_papers_HM/dev_nemo_format.tsv"

# 사전 학습된 모델 이름 설정
PRETRAINED_MODEL_NAME = 'bert-base-uncased'

# 학습률 설정
LR = 1e-4

# 설정 값 평가를 위해 저장
with open("my_assessment/step2.json", "w") as outfile: 
    json.dump([MAX_SEQ_LENGTH, NUM_CLASSES, BATCH_SIZE], outfile)

# 트레이너와 exp_manager 설정 출력
print(OmegaConf.to_yaml(config.trainer))
print(OmegaConf.to_yaml(config.exp_manager)) 

# 추가 설정 값 정의
MAX_EPOCHS = 5  # 최대 에포크 수
AMP_LEVEL = 'O1'  # 혼합 정밀도 설정
PRECISION = 16  # 16비트 정밀도 사용

# 설정 값 평가를 위해 저장
with open("my_assessment/step3.json", "w") as outfile: 
    json.dump([MAX_EPOCHS, AMP_LEVEL, PRECISION], outfile) 

# 모델 학습 스크립트 실행, 명령줄에서 설정 값을 재정의
TC_DIR = "/dli/task/nemo/examples/nlp/text_classification"

os.system(f'''
python {TC_DIR}/text_classification_with_bert.py \
    model.dataset.num_classes={NUM_CLASSES} \
    model.dataset.max_seq_length={MAX_SEQ_LENGTH} \
    model.train_ds.file_path={PATH_TO_TRAIN_FILE} \
    model.validation_ds.file_path={PATH_TO_VAL_FILE} \
    trainer.max_epochs={MAX_EPOCHS} \
    model.language_model.pretrained_model_name={PRETRAINED_MODEL_NAME} \
    trainer.gpus=1 \
    trainer.precision={PRECISION} \
    trainer.amp_level={AMP_LEVEL} \
    exp_manager.exp_dir="/dli/task/experiments" \
    model.optim.lr={LR}
''')

# 학습 후 명령어와 로그를 평가를 위해 저장
cmd_log = os.path.join(os.path.dirname(os.path.dirname(get_latest_model())),'cmd-args.log')
lightning_logs = os.path.join(os.path.dirname(os.path.dirname(get_latest_model())),'lightning_logs.txt')

with open(cmd_log, "r") as f:
    cmd = f.read()
    cmd_list = cmd.split()
with open("my_assessment/step4.json", "w") as outfile: 
    json.dump(cmd_list, outfile)
    
with open(lightning_logs, "r") as f:
    log = f.readlines()
with open("my_assessment/step4_lightning.json", "w") as outfile:
    json.dump(log, outfile)

# 최신 모델 체크포인트에서 모델 복원
model = nemo_nlp.models.TextClassificationModel.restore_from(get_latest_model())

# 테스트 파일 목록 설정
test_files = [
    'test49.tsv',
    'test50.tsv',
    'test51.tsv',
    'test52.tsv',
    'test53.tsv',
    'test54.tsv', 
    'test55.tsv',
    'test56.tsv',
    'test57.tsv',
    'test62.tsv',
]

# 각 테스트 파일에 대해 예측 수행
results = []
for test_file in test_files:
    filepath = os.path.join(DATA_DIR, test_file)
    with open(filepath, "r") as f:
        lines = f.readlines()
    del lines[0]
    
    results.append(model.classifytext(lines, batch_size=1, max_seq_length=256))

print(results)

# 결과를 기반으로 저자를 예측하고 저장
author = []
for result in results:
    avg_result = sum(result) / len(result)
    if avg_result < 0.5:
        author.append("HAMILTON")
        print("HAMILTON")
    else:
        author.append("MADISON")
        print("MADISON")
        
with open("my_assessment/step5.json", "w") as outfile: 
    json.dump(author, outfile)