DEVOCEAN

!pip install transformers datasets bitsandbytes peft trl accelerate --upgrade -qqq
!pip install -U huggingface_hub
!pip3 install -q -U datasets==2.18.0

import torch
from datasets import Dataset, load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline, TrainingArguments
from peft import LoraConfig, PeftModel
from trl import SFTTrainer

from huggingface_hub import notebook_login
notebook_login()

dataset = load_dataset("daekeun-ml/naver-news-summarization-ko")

Dataset({
    features: ['date', 'category', 'press', 'title', 'document', 'link', 'summary'],
    num_rows: 22194
})

{'date': '2022-07-03 17:14:37',
 'category': 'economy',
 'press': 'YTN ',
 'title': '추경호 중기 수출지원 총력 무역금융 40조 확대',
 'document': '앵커 정부가 올해 하반기 우리 경제의 버팀목인 수출 확대를 위해 총력을 기울이기로 했습니다. 특히 수출 중소기업의 물류난 해소를 위해 무역금융 규모를 40조 원 이상 확대하고 물류비 지원과 임시선박 투입 등을 추진하기로 했습니다. 류환홍 기자가 보도합니다. 기자 수출은 최고의 실적을 보였지만 수입액이 급증하면서 올해 상반기 우리나라 무역수지는 역대 최악인 103억 달러 적자를 기록했습니다. 정부가 수출확대에 총력을 기울이기로 한 것은 원자재 가격 상승 등 대외 리스크가 가중되는 상황에서 수출 증가세 지속이야말로 한국경제의 회복을 위한 열쇠라고 본 것입니다. 추경호 경제부총리 겸 기획재정부 장관 정부는 우리 경제의 성장엔진인 수출이 높은 증가세를 지속할 수 있도록 총력을 다하겠습니다. 우선 물류 부담 증가 원자재 가격 상승 등 가중되고 있는 대외 리스크에 대해 적극 대응하겠습니다. 특히 중소기업과 중견기업 수출 지원을 위해 무역금융 규모를 연초 목표보다 40조 원 늘린 301조 원까지 확대하고 물류비 부담을 줄이기 위한 대책도 마련했습니다. 이창양 산업통상자원부 장관 국제 해상운임이 안정될 때까지 월 4척 이상의 임시선박을 지속 투입하는 한편 중소기업 전용 선복 적재 용량 도 현재보다 주당 50TEU 늘려 공급하겠습니다. 하반기에 우리 기업들의 수출 기회를 늘리기 위해 2 500여 개 수출기업을 대상으로 해외 전시회 참가를 지원하는 등 마케팅 지원도 벌이기로 했습니다. 정부는 또 이달 중으로 반도체를 비롯한 첨단 산업 육성 전략을 마련해 수출 증가세를 뒷받침하고 에너지 소비를 줄이기 위한 효율화 방안을 마련해 무역수지 개선에 나서기로 했습니다. YTN 류환홍입니다.',
 'link': 'https://n.news.naver.com/mnews/article/052/0001759333?sid=101',
 'summary': '올해 상반기 우리나라 무역수지는 역대 최악인 103억 달러 적자를 기록한 가운데, 정부가 하반기에 우리 경제의 버팀목인 수출 확대를 위해 총력을 기울이기로 결정한 가운데, 특히 수출 중소기업의 물류난 해소를 위해 무역금융 규모를 40조 원 이상 확대하고 물류비 지원과 임시선박 투입 등을 추진하기로 했다.'}

BASE_MODEL = "meta-llama/Llama-2-7b-hf"

lora_config = LoraConfig(
    r=8,
    lora_alpha = 8,
    lora_dropout = 0.05,
    target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
    bias="none",
    task_type="CAUSAL_LM",
)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(BASE_MODEL,
                                             quantization_config=bnb_config,
                                             device_map="auto"
                                            )

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset['train'],
    max_seq_length=512,
    args=TrainingArguments(
        output_dir="outputs",
#        num_train_epochs = 1,
        max_steps=500,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        optim="paged_adamw_8bit",
        warmup_steps=0.03,
        learning_rate=2e-4,
        fp16=True,
        logging_steps=100,
        push_to_hub=False,
        report_to='none',
    ),
    peft_config=lora_config,
    formatting_func=generate_prompts,
)

trainer.train()

TrainOutput(global_step=500, 
training_loss=0.9232814483642579, 
metrics = {
'train_runtime': 3351.6021, 
'train_samples_per_second': 0.597, 
'train_steps_per_second': 0.149, 
'total_flos': 4.038108997678694e+16, 
'train_loss': 0.9232814483642579, 
'epoch': 0.09011444534558889})

ADAPTER_MODEL = "lora_adapter"

trainer.model.save_pretrained(ADAPTER_MODEL)

model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, device_map='auto', torch_dtype=torch.float16)
model = PeftModel.from_pretrained(model, ADAPTER_MODEL, device_map='auto', torch_dtype=torch.float16)

model = model.merge_and_unload()
model.save_pretrained('./Llama-2-7b-hf-finetuned')

BASE_MODEL = "meta-llama/Llama-2-7b-hf"
FT_MODEL = "./Llama-2-7b-hf-finetuned"

lora_config = LoraConfig(
    r=8,
    lora_alpha = 8,
    lora_dropout = 0.05,
    target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
    bias="none",
    task_type="CAUSAL_LM",
)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

ft_model = AutoModelForCausalLM.from_pretrained(FT_MODEL,
quantization_config=bnb_config, device_map={"":0})
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, add_special_tokens=True)

def generate_prompts_sentence(example):
    prompt = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>다음 글을 요약해주세요:
{example['document']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
{example['summary']}<|eot_id|>"""
    return prompt

pipe_finetuned = pipeline("text-generation", model=model2, tokenizer=tokenizer, max_new_tokens=512)

outputs = pipe_finetuned(
    generate_prompts_sentence(dataset['test'][number]),
    do_sample=True,
    temperature=0.2,
    top_k=50,
    top_p=0.95,
    add_special_tokens=True
)

print(outputs[0]['generated_text'][len(generate_prompts_sentence(dataset['test'][number])):])