code_test - taewook

'''
---taewook/code_test Auto Generate Code---
Author		: taewook
Project Name: code_test
Project Link: https://blockai.kr/taewook/code_test (BlockAI)
Create Date	: 2024-08-23

---Requirements---
# 사용자의 환경(OS, CUDA 등)에 따라 라이브러리 버전을 맞춰주세요
pip install torch==2.0 torchvision==0.15.2 torchtext==0.15.2 torchaudio==2.0.2
pip install pytorch-lightning==2.0.4
pip install tqdm
pip install pandas
pip install scikit-learn
pip install transformers
pip install timm

---Folder Structure---
--📂 data
   |--📄 train.csv
   |--📄 test.csv
--📄 code_test.py
--📄 code_test.ipynb
--📄 requirements.txt
'''

import os
import argparse
import copy
from glob import glob

from tqdm import tqdm
import numpy as np
import pandas as pd
from sklearn import preprocessing
from sklearn.model_selection import train_test_split

import torch
import pytorch_lightning as pl

path_sep = os.sep


# https://pytorch.org/tutorials/beginner/basics/data_tutorial.html#creating-a-custom-dataset-for-your-files
class Dataset(torch.utils.data.Dataset):
    def __init__(self, inputs, targets=[]):
        self.inputs = inputs
        self.targets = targets

    # 학습 및 추론 과정에서 데이터를 1개씩 꺼내오는 곳
    def __getitem__(self, idx):
        # 정답이 있다면 if문을, 없다면 else문을 수행합니다
        if len(self.targets) == 0:
            return torch.tensor(self.inputs[idx])
        else:
            return torch.tensor(self.inputs[idx]), torch.tensor(self.targets[idx])

    # 입력하는 개수만큼 데이터를 사용합니다
    # 'return 100'이면 1에폭에 100개의 데이터만 사용합니다
    def __len__(self):
        return len(self.inputs)


# https://pytorch-lightning.readthedocs.io/en/stable/extensions/datamodules.html
class Dataloader(pl.LightningDataModule):
# 데이터의 종류에 따라 코드 수정이 필요할 수 있습니다
    def __init__(self, data_folder, batch_size, train_ratio, shuffle):
        super().__init__()
        self.data_folder = data_folder
        self.batch_size = batch_size
        self.train_ratio = train_ratio
        self.shuffle = shuffle

        self.train_dataset = None
        self.test_dataset = None
        self.predict_dataset = None

    def preprocessing(self, data):
        # 타겟 데이터가 없으면 빈 배열을 리턴합니다.
        try:
            targets = data[self.target_columns].values.tolist()
            inputs = data.drop(self.target_columns, axis=1).values.tolist()
        except:
            targets = []
            inputs = data.values.tolist()

        return inputs, targets

    def setup(self, stage='fit'):
        if stage == 'fit':
            train_data = pd.read_csv(os.path.join(self.data_folder, 'train.csv'), sep=',', header=None, index_col=None, encoding='utf-8')

            # 학습데이터 준비
            train_inputs, train_targets = self.preprocessing(train_data)

            # train 데이터만 shuffle을 적용해줍니다, 필요하다면 val, test 데이터에도 shuffle을 적용할 수 있습니다
            self.train_dataset = Dataset(train_inputs, train_targets)
        else:
            # 평가데이터 준비
            test_data = pd.read_csv(os.path.join(self.data_folder, 'test.csv'), sep=',', header=None, index_col=None, encoding='utf-8')
            test_inputs, test_targets = self.preprocessing(test_data)
            self.test_dataset = Dataset(test_inputs, test_targets)
            self.predict_dataset = Dataset(test_inputs, [])

    def train_dataloader(self):
        return torch.utils.data.DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=args.shuffle)

    def test_dataloader(self):
        return torch.utils.data.DataLoader(self.test_dataset, batch_size=self.batch_size)

    def predict_dataloader(self):
        return torch.utils.data.DataLoader(self.predict_dataset, batch_size=self.batch_size)


# https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html
class Model(pl.LightningModule):

if __name__ == '__main__':
    # https://docs.python.org/ko/3/library/argparse.html
    # 하이퍼 파라미터 등 각종 설정값을 입력받습니다
    # 터미널 실행 예시 : python3 run.py --batch_size=64 ...
    # 실행 시 '--batch_size=64' 같은 인자를 입력하지 않으면 default 값이 기본으로 실행됩니다
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_folder', default='./data')
    parser.add_argument('--batch_size', default=0)
    parser.add_argument('--max_epoch', default=0)
    parser.add_argument('--shuffle', default=False)
    parser.add_argument('--train_ratio', default=1.0)
    args = parser.parse_args()

    dataloader = Dataloader(args.data_folder, args.batch_size, args.train_ratio, args.shuffle)
    model = Model()
    # https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html
    # 학습 및 추론을 위한 Trainer 설정
    trainer = pl.Trainer(accelerator='gpu', devices=1, max_epochs=args.max_epoch)
    trainer.fit(model=model, datamodule=dataloader)
    # trainer.test(model=model, datamodule=dataloader)
    # predictions = trainer.predict(model=model, datamodule=dataloader)