Machine Learning/NLP

[code]BERT example code

뚜둔뚜둔 2023. 7. 20. 00:13
import torch
import transformers import BertTokenizerFast
from torch.utils.data import Dataset, DataLoader

class TokenDataset(Dataset):

	def __init__(self, dataframe, tokenoizer_pretrained):
    	# sentence, label 컬럼으로 구성된 데이터프레임 전달
    	self.data = dataframe
        # huggingface 토크나이저 생성
        self.tokenizer = BertTokenizerFast.from_pretrained(tokenizer_pretrained)

	def __len__(self):
    	return len(self.data)
        
    def __getitem__(self, idx):
    	sentence = self.data.iloc[idx]['document']
        label=self.data.iloc[idx]['label']
        
        # tokenizer 처리
        tokens = self.tokenizer(
        	sentence,					# 1개 문장
            return_tensors='pt',		# 텐서로 변환
            truncation=Ture,			# 잘라내기 허용
            padding='max_length',		# 패딩 적용
            add_special_tokens=True		#스페셜 토큰 적용
        )
        
        input_idx=tokens['input_ids'].squeeze(0) 				# 2D -> 1D
        attention_mask =tokens['attention_mask'].squeeze(0)		# 2D -> 1D
        token_type_idx = torch.zeros_like(attention_mask)
        
        # input_ids, attention_mask, token_type_idx 이렇게 3가지 요소를 반환하도록 함
        # input_idx : token
        # attention_mask: 실제 단어가 존재하면1, padding이면 0 (패딩은 0이 아닐 수 있음)
        # token_type_ids: 문장을 구분하는 id, 단일 문장인 경우에는 전부 0) 
        
        return {
        	'input_ids':input_ids,
            'attention_mask':attention_mask,
            'token_type_idx': token_type_ids,
        }, torch.tensor(label)

 

nn.Dataset을 확장한 클래스인 TokenDataset 클래스 인스턴스를 생성한다

 

 

 

 

 

 

 

 

 

 

 

 

 

 

참고 (https://teddylee777.github.io/huggingface/bert-kor-text-classification/)

내용을 참고 하여 작성하였습니다

코드 전체를 따라서 돌려보시려면 참고 사이트에 가시면 전체 원본이 있습니다.

반응형