[NLP] EXAONE 모델 가져다가 추론만 시키기

멘토링때 멘토님이 EXAONE 모델도 추천을 해주셔서 대회 하루 남겨놓고 급하게 추론만이라도 돌려본다. 30분이 지났는데도 아직도 요약문 만드는 중.. 생각보다 오래걸린다. 

현재 사용하고 있는 모델: BM-K/EXAONE-3.0-7.8B-Daily-Conversation-Summary

링크에 가면 Quick Start라고 어떻게 쓰면 되는지 나와있는데, 이 방법대로 따라갔다.

먼저 모델을 불러와준다. (트랜스포머 버전 안맞으면 로드가 안됨.. 관련글)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

device = torch.device('cuda:0' if torch.cuda.is_available()  else 'cpu')

model_name = "BM-K/EXAONE-3.0-7.8B-Daily-Conversation-Summary"
model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        load_in_8bit=True, 
        trust_remote_code=True,
        low_cpu_mem_usage=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

여기서 또 모델이 커서 그런건지 첨에 로드가 잘 안되서 

load_in_8bit=True 라는 인자를 줬다. 이거 주려면 아래 인스톨 해야함.. 이게 뭐 모델을 손쉽게 압축할 수 있다나? 4비트 양자화 8비트 양자화 어쩌고 저쩌고 하는데, 아 이부분은 진짜 하나도 모르겠다. 일단 패스.. 

pip install -U bitsandbytes

이러고 Loading checkpoint shards 100%가 뜨면 모델이 로드된 것!

그리고 프롬프트와 대화문을 줘야 하는데, Quick Start에서는 대화문이 그냥 한 줄이였다.

PROMPT = ''' '''
chat = " "

내가 가진 Jsonl 파일에는 딕셔너리 리스트 형태로 대화문이 저장되어 있다. 구조는 아래처럼 되어 있고, 지금 딱 한 줄만 가져온거.. id가 500개인 json파일. (능력자 팀원분이 만들어주신 것)

{
    "id": "nikluge-2024-일상 대화 요약-test_0",
    "input": {
        "conversation": [
            {"speaker": "Person1", "utterance": "더슨 씨, 받아쓰기 좀 해주세요."},
            {"speaker": "Person2", "utterance": "네, 실장님..."},
            {"speaker": "Person1", "utterance": "이것은 오늘 오후까지 모든 직원에게 내부 메모로 전달되어야 합니다. 준비되셨나요?"}
        ]
    }
}

여기서 id는 필요없고, conversation에서 person1: 더슨씨... 이 내용만 뽑아서 메세지에 담아주려고 했다.

#이 형태로 만들려고 하는 것!! 

message = [
        {"role": "system", "content": PROMPT},
        {"role": "user", "content": chat}
]

그래서 json에서 내용만 뽑아서 리턴하는 함수를 먼저 만들어주고, 그걸로 chat_list를 만들어준다. 

import json

def process_jsonl(file_path):
    chat = []
    
    with open(file_path, 'r', encoding='utf-8') as file:
        for line in file:
            data = json.loads(line.strip())
            conversation = data['input']['conversation']
            formatted_conv = '\n'.join(
                f"{turn['speaker']}: {turn['utterance']}"
                for turn in conversation)
            
            content = f"""{formatted_conv}"""
            
            chat.append({
                'id': data['id'],
                'content': content
            })
    
    return chat

file_path = '/data/ephemeral/home/exp9/exaone_test.jsonl'
chat_list = process_jsonl(file_path)


chat_list[1]['content']
'''
'Person1: 드디어 왔네! 왜 그렇게 오래 걸렸어?\nPerson2: 또 교통 체증에 걸렸어. ...
'''

chat_list에 잘 저장이 된 걸 확인하고, 이걸 메세지에 프롬프트와 함께 넣어준다.

PROMPT = '''
You are a summarization expert of the XAONE model from LG AI Research. 
Please read the following conversation and summarize it. 
When using the words #Person1# and #Person2#, 
please keep the # symbols as they are in the summary.
'''

message = []

for i in range(len(chat_list)):
    message.append([
        {"role": "system", "content": PROMPT},
        {"role": "user", "content": chat_list[i]['content']}
        ])
        
        
message[0]

'''
[{'role': 'system',
  'content': 'You are a summarization expert of the XAONE model ... '},
 {'role': 'user',
  'content': 'Person1: 더슨 씨, 받아쓰기 좀 해주세요.\nPerson2: 네, 실장님...'}]
'''

메세지에 잘 담겼다. 이제 이 메세지를 모델이 이해할 수 있도록 형식을 변환한다. PyTorch 텐서 형태로 반환한다. 

source = tokenizer.apply_chat_template(
        message,
        add_generation_prompt=True,
        padding=True,
        truncation=True,
        return_tensors="pt"
)

메세지가 들어간 뒤에 출력을 얻는다. 

from tqdm import tqdm

model.gradient_checkpointing_enable()

results = []

for i in tqdm(range(len(message)), desc="Generating summaries"):
    source = tokenizer.apply_chat_template(
        message[i],
        add_generation_prompt=True,
        padding=True,
        truncation=True,
        max_length=512,
        return_tensors="pt"
    )
    
    output = model.generate(
        source.to(device),
        max_new_tokens=64,
        eos_token_id=tokenizer.eos_token_id,
        do_sample=False
    )
    
    #decoded_output = tokenizer.decode(output[0], skip_special_tokens=True
    decoded_output = tokenizer.decode(output[0][source.shape[-1]:], skip_special_tokens=True)
    
    results.append({
        'id': chat_list[i]['id'],
        'summary': decoded_output
    })

    torch.cuda.empty_cache()
    gc.collect()

이거 첨에 모르고.. 

decoded_output = tokenizer.decode(output[0], skip_special_tokens=True 로 설정해서 응답값에 이런 식으로 모두 다 들어갔다ㅠㅠ [|assistant|] 값만 빼려고 했더니 잘 안되서 프롬프트에 퓨샷 추가하고, 맥스 길이 늘려서 다시 돌려보는 중..

'[|system|]You are a summarization expert of the XAONE ...
[|user|]Person1: 더슨 씨, 받아쓰기 좀 해주세요.\nPerson2: 네, 실장님...
[|assistant|]이 대화에서 #Person1#은 직원들에게 새로운 통신 정책을 전달해달라고 요청했습니다. #Person2#는 이 정책이 내부뿐만 아니라 외부 통신에도 적용된다고 확인했습니다. #Person1#은 즉시 메시지 사용을 금지하고, 이를 위반할 경우 경고'

모델 자체를 파인튜닝하고 싶어서 다른 팀원이 시도해봤는데, vram이 모자라서 안된다고 한다.

흠 파인튜닝하면 괜찮을거같은데 아쉽구만.

 

 

 

(+ 내가 베이스라인 코드 보면서 추론 부분만 떼서 만들었던 코드)

import torch
from tqdm import tqdm

from transformers import AutoTokenizer, AutoModelForCausalLM

device = torch.device('cuda:0' if torch.cuda.is_available()  else 'cpu')

model_name = "BM-K/EXAONE-3.0-7.8B-Daily-Conversation-Summary"
model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map=device,
        trust_remote_code=True,
        low_cpu_mem_usage=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name)


max_length = 128
# 전처리 함수 정의 (입력 텍스트만 처리)
def preprocess_test_data(batch):
    processed_inputs = []
    for conversation in batch["input"]:
        utterances = [f"{item['speaker']}: {item['utterance']}" for item in conversation]
        processed_inputs.append(" ".join(utterances))
    
    batch["input"] = processed_inputs


    inputs = tokenizer(
        batch["input"], 
        max_length=max_length, 
        truncation=True, 
        padding="max_length"
    )
    batch["input_ids"] = inputs["input_ids"]
    batch["attention_mask"] = inputs["attention_mask"]
    return batch

# 후처리 함수 정의
def clean_generated_summary(summary, tokenizer):
    # 모델의 기본 특수 토큰 제거
    if tokenizer.bos_token:
        summary = summary.replace(tokenizer.bos_token, "")
    if tokenizer.eos_token:
        summary = summary.replace(tokenizer.eos_token, "")
    if tokenizer.pad_token:
        summary = summary.replace(tokenizer.pad_token, "")
    # 불필요한 공백 제거
    summary = summary.strip()
    return summary

# 추론 함수 정의
def generate_summaries(model, dataloader, tokenizer, device):
    model.eval()
    summaries = []
    
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="Generating summaries"):
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            
            # 요약 생성
            outputs = model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=100,
                num_beams=4,  # 빔 서치의 빔 수 (성능과 속도의 균형)
                early_stopping=True,
                no_repeat_ngram_size=2  # 반복 방지
            )
            
            # 디코딩 (특수 토큰 포함)
            decoded_summaries = tokenizer.batch_decode(outputs, skip_special_tokens=False)
            
            # 후처리: 불필요한 특수 토큰 제거
            cleaned_summaries = [clean_generated_summary(summary, tokenizer) for summary in decoded_summaries]
            summaries.extend(cleaned_summaries)
    
    return summaries

def extract_input(batch):
    conversations = []
    for example in batch["input"]:
        conversations.append(example["conversation"])
    return {"input": conversations}

    # # JSON에서 "input" 데이터를 적절히 추출
    # return {"input": batch["input"]["conversation"]}
    


from datasets import load_dataset, Dataset, concatenate_datasets
from torch.utils.data import DataLoader
import pandas as pd


special_tokens_dict = {
    "additional_special_tokens": [
        '#Person1#', '#Person2#', '#Person3#', '#Person4#', '#Person5#', '#Person6#', '#Person7#',
        '#PhoneNumber#', '#Address#', '#PassportNumber#', '#DateOfBirth#', 
        '#CarNumber#', '#Email#', '#SSN#', '#CardNumber#'
    ]
}
tokenizer.add_special_tokens(special_tokens_dict)
model.resize_token_embeddings(len(tokenizer))

# 모델을 평가 모드로 전환
model.eval()

# GPU 사용 가능 시 GPU로 이동
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 테스트 데이터 로드
test_file = "/data/ephemeral/home/exp9/exaone_test.jsonl"
# jsonl 파일을 Dataset 객체로 로드
test_dataset = Dataset.from_json(test_file)

# 전처리 단계
test_dataset = test_dataset.map(extract_input, batched=True)

# 데이터 전처리
tokenized_test = test_dataset.map(preprocess_test_data, batched=True)


# 데이터셋을 PyTorch 텐서 형식으로 설정
tokenized_test.set_format(type='torch', columns=['input_ids', 'attention_mask'])

# DataLoader 설정
batch_size = 16  # GPU 메모리에 맞게 조절
test_loader = DataLoader(tokenized_test, batch_size=batch_size)

# 요약 생성
summaries = generate_summaries(model, test_loader, tokenizer, device)

# 원본 테스트 데이터에 요약 추가
test_dataset = test_dataset.add_column("summary", summaries)

# 결과 저장
csv_output_path = "/data/ephemeral/home/exp9/exaone_test_with_summaries.csv"
test_dataset.to_csv(csv_output_path)
print(f"추론 결과가 저장되었습니다.")

이렇게 만든 결과.. 

'Person1: 더슨 씨, 받아쓰기 좀 해주세요. Person2: 네, 실장님... Person1: 이것은 오늘 오후까지 모든 직원에게 내부 메모로 전달되어야 합니다. 준비되셨나요? Person2: 네, 실장님. 시작하셔도 됩니다. Person1: 모든 직원들에게 주의하라... 즉시 효력을 발휘하여, 모든 사무실 통신은 이메일 통신과 공식 메모로 제한됩니다. 근무 시간 동안 직원들이 즉시 메시지 프로그램을 사용하는 것은 엄격히 금지됩니다. Person2: 실장님, 이것은 내부 통신에만 적용되는 건가요? 아니면 외부 통신에도 해당되는지요?

\n\n이 대화에서 두 사람은 다음과 같은 주제에 대해 이야기했습니다:

\n- 회사의 새로운 정책에 대한 전달 및 시행\n  
- 새로운 통신 정책의 시행과 관련된 대화가 이루어졌습니다. 이 정책은 회사 내부의 모든 통신을 이메일과 메모 등 공식적인 방법으로만 하도록 규정하고 있으며, 비공식적 통신 수단인 메신저 프로그램의 사용을 금지합니다. 이러한 대화를 통해'

실제 테스트파일 보니까 들어가는 대화문이 짤렸고, 

뒷부분 요약하는 부분만 나오면 되는데 이렇게 해서는 내가 따로 조절할 수가 없는 것 같아서.. Quick start 대로 프롬프트 주고 답받는 방식으로 돌아감!