Chrom256/gemma-2-9b-it-lora_20241216_033631

Uploaded model

• Developed by: Chrom256
• License: gemma
• Finetuned from model : google/gemma-2-9b

This gemma2 model was trained 2x faster with Unsloth and Huggingface's TRL library.

決勝戦進出後、本READMEに解説等を追記しましたが、実行コード自体に変更はありません。

概要

GoogleのGemma-2-9bをベースに、LoRAを使用して最適化した日本語対応モデルです。

本モデルは東京大学松尾・岩澤研究室大 規模言語モデル Deep Learning 応用講座(2024)におけるLLMコンペティション提出用に作成されたものです。

訓練データ

以下のデータセットを用いてInstruction-tuningを実施しました。 Googleより公開されている、Gemma用の指示チューニング用フォーマットを参考にしました。

• llm-jp/magpie-sft-v1.0 (Apache License 2.0)

全て使用しました

• Aratako/Magpie-Tanuki-8B-annotated-96k (Apache License 2.0)

primary_tagでCoding & Debuggingに該当するものは除外し、その上でinput_qualityがgood、excellentに該当するものを使用しました

使用方法

このモデルは、以下のGoogle Colabノートブックで実行できます

Google Colabで実行

環境構築が容易である事から、上記リンクからGoogle Colabで実行する形を推奨します。 Google Colab以外の環境で実行される場合には、Google Colabのリンク先からコードをダウンロードし、環境に合わせて修正を加えた上でご使用ください。

Google Colabでの実行前準備

• シークレットにHF_TOKENを配置してください
  • 名前に'HF_TOKEN'を入力
  • 値にHugging Faceトークンを入力して保存
• ファイルタブ内にelyza-tasks-100-TV_0.jsonlを配置してください

L4でのコード実行時間は全体で約45分でした。

• 事前準備（環境構築、モデルのロード等）：約18分
• 推論：約27分

メモリが不足する事はありませんでした。 出力物は新規作成されるoutputファイルの中に、Output.jsonl、として出力されます。

通信環境によっては出力されたOutput.jsonlファイルが見えないことがありますが、ファイルタブ内での再読み込み等により表示、ダウンロードできる事を複数回の検証により確認しています。

ライセンス

本モデルは利用したモデル、訓練データの関係で以下のライセンスの影響を受けます。

• Gemma Terms of Use
• 使用した訓練データのライセンス

導入・推論用コードの解説

正確なコードは先ほど提示したGoogle Colabのリンク先をご参照ください。Google Colab上で最後まで問題なく実行できる事を確認しております。

Google Colab以外の環境で実行される場合には、リンク先からコードをダウンロードし、環境に合わせ修正を加えた上でご使用ください。ダウンロードができない方は、本READMEの「実行用コード全文」のコードをご利用ください。

＊本READMEはコードをコピペし編集を加えています。そのため、READMEからコピペしてローカルで実行すると、実行エラーが起きる可能性が高いため、非推奨です。

以下、コードの解説です。

環境構築部分

Google Colab以外で実行する際には、HF_TOKENは直接コード中に記載してください。 現在はモデルがPublicとなっているため、HF_TOKENは必要ないかもしれません。 しかし、コンペに提出したJSONLファイルを出力する際に使用したコードと、全く同じコードを提出するのが望ましいと判断し、記述を残したままにしています。

!pip install -q transformers==4.46.3 accelerate bitsandbytes
!pip install -q tqdm
!pip install flash-attn --no-build-isolation

import os
import torch
import json
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from torch.cuda.amp import autocast
from concurrent.futures import ThreadPoolExecutor
import threading

#モデルをロードするのに必要なHF_TOKENの読み込み
from google.colab import userdata
HF_TOKEN = userdata.get('HF_TOKEN')

if HF_TOKEN is None:
    raise ValueError("HF_TOKENが設定されていません。上記の手順でトークンを設定してください。")

モデルの準備部分

ロードを高速化する目的でPythonのThreadPoolExecutorクラスを使用し、マルチスレッド処理を実装しました。

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

#モデルのロード
def load_model_and_tokenizer():
    model_id = "Chrom256/gemma-2-9b-it-lora_20241216_033631"
    base_model_id = "google/gemma-2-9b"
    downloaded_components = {"model": None, "tokenizer": None}
    download_lock = threading.Lock()

    def download_base_model():
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.bfloat16
        )

        model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            quantization_config=quantization_config,
            device_map="auto",
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            attn_implementation="eager",
            low_cpu_mem_usage=True,
            token=HF_TOKEN
        )
        with download_lock:
            downloaded_components["model"] = model

    def download_tokenizer():
        tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            trust_remote_code=True,
            token=HF_TOKEN
        )
        with download_lock:
            downloaded_components["tokenizer"] = tokenizer


    torch.cuda.empty_cache()

    #ロードを高速化する目的でThreadPoolExecutorを使用しました
    with ThreadPoolExecutor(max_workers=2) as executor:
        model_future = executor.submit(download_base_model)
        tokenizer_future = executor.submit(download_tokenizer)
        model_future.result()
        tokenizer_future.result()

    model = downloaded_components["model"]
    tokenizer = downloaded_components["tokenizer"]
    
    #GPUのキャッシュをクリアする事で、メモリが足りなくなる状況を防ぎます
    torch.cuda.empty_cache()
    
    #エラー発生対策
    try:
        adapter_path = model_id
        print(f"Loading adapter from {adapter_path}")
        model.load_adapter(adapter_path, "default", token=HF_TOKEN)
        print("Adapter loaded successfully")
    except Exception as e:
        print(f"Error loading adapter: {e}")
        raise


    model.config.use_cache = True
    model.eval()

    torch.cuda.empty_cache()

    return model, tokenizer

推論部分

簡潔で装飾や特殊記号を含まない出力が出るようプロンプトを設定すると共に、コンペルール内で出力の整形を行っています。

def run_inference(model, tokenizer, tokenized_inputs, generation_config, batch_size=4):
    results = []

    for i in tqdm(range(0, len(tokenized_inputs), batch_size)):
        batch = tokenized_inputs[i:i+batch_size]

        prompts = [
            f"""<start_of_turn>system
簡潔に回答してください。装飾や特殊記号は使用しないでください。
<end_of_turn>
<start_of_turn>user
{item["input"]}
<end_of_turn>
<start_of_turn>model
""" for item in batch
        ]

        inputs = tokenizer(
            prompts,
            padding=True,
            truncation=True,
            return_tensors="pt"
        ).to(model.device)

        with torch.no_grad(), autocast(dtype=torch.bfloat16):
            outputs = model.generate(
                **inputs,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                **generation_config
            )

            for idx, output in enumerate(outputs):
                response = tokenizer.decode(output, skip_special_tokens=True)

                if 'model\n' in response:
                    response = response.split('model\n')[-1].strip()
                elif 'model' in response:
                    response = response.split('model')[-1].strip()


                response = post_process_output(response)

                results.append({
                    "task_id": batch[idx]["task_id"],
                    "input": batch[idx]["input"],
                    "output": response
                })


        del outputs, inputs
        torch.cuda.empty_cache()

    return results

#出力には余計な装飾が含まれている事があったため、コンペルールに従った方法で出力を整形
def post_process_output(response):
    response = response.strip()
    symbols_to_replace = ['**', '`', '|', '```', '---', '===']
    for symbol in symbols_to_replace:
        response = response.replace(symbol, ' ')
    return ' '.join(response.split())

#設定
GENERATION_CONFIG = {
    "max_new_tokens": 512,
    "use_cache": True,
    "do_sample": False,
    "num_beams": 4,
    "repetition_penalty": 1.2,
    "length_penalty": 1.0,
    "early_stopping": False
}

他の関数

問題ファイルを読み込む関数と、出力をJSONL形式で出力する関数です。

#問題の読み込み
def load_input_data(file_path):
    tokenized_inputs = []
    with open(file_path, "r") as f:
        for line in f:
            if line.strip():
                dt = json.loads(line)
                tokenized_inputs.append({
                    "task_id": dt["task_id"],
                    "input": dt["input"]
                })
    return tokenized_inputs

#出力の保存
def save_results(results, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    jsonl_path = os.path.join(output_dir, "Output.jsonl")

    with open(jsonl_path, 'w', encoding='utf-8') as f:
        for item in results:
            json.dump(item, f, ensure_ascii=False)
            f.write('\n')

    print(f"Saved results to: {jsonl_path}")

全て実行

以上の関数を順番に実行

def main():
    #モデルのロード
    model, tokenizer = load_model_and_tokenizer()
    #推論する問題の読み込み
    tokenized_inputs = load_input_data("/content/elyza-tasks-100-TV_0.jsonl")
    #推論実行
    results = run_inference(model, tokenizer, tokenized_inputs, GENERATION_CONFIG)
    #結果を保存
    save_results(results, "output")

if __name__ == "__main__":
    main()

実行用コード全文

念の為、推論用のコード全文を提示します。 提示したGoogle Colabリンク先よりダウンロードしたPythonコードをコピペしたものです。 Google Colabよりコードがダウンロードが出来ない場合にご利用ください。

# -*- coding: utf-8 -*-
"""共有中：推論コード.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/10MWpy05Xw7UT5ySM2QpOdb2aKqwRjj8m
"""

!pip install -q transformers==4.46.3 accelerate bitsandbytes
!pip install -q tqdm
!pip install flash-attn --no-build-isolation

import os
import torch
import json
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from torch.cuda.amp import autocast
from concurrent.futures import ThreadPoolExecutor
import threading

print("【重要】以下の手順でHugging Faceトークンを設定しておいてください")
print("1. 左メニューの'シークレット'タブを開く")
print("2. '新しいシークレット'をクリック")
print("3. 名前に'HF_TOKEN'を入力")
print("4. 値にHugging Faceトークンを入力して保存")
print("ファイルタブ内にelyza-tasks-100-TV_0.jsonlを配置しておいてください")
print("出力物は、新規に作成されるoutputファイルの中に格納されます")

from google.colab import userdata
HF_TOKEN = userdata.get('HF_TOKEN')

if HF_TOKEN is None:
    raise ValueError("HF_TOKENが設定されていません。上記の手順でトークンを設定してください。")

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

def load_model_and_tokenizer():
    model_id = "Chrom256/gemma-2-9b-it-lora_20241216_033631"
    base_model_id = "google/gemma-2-9b"
    downloaded_components = {"model": None, "tokenizer": None}
    download_lock = threading.Lock()

    def download_base_model():
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.bfloat16
        )

        model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            quantization_config=quantization_config,
            device_map="auto",
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            attn_implementation="eager",
            low_cpu_mem_usage=True,
            token=HF_TOKEN
        )
        with download_lock:
            downloaded_components["model"] = model

    def download_tokenizer():
        tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            trust_remote_code=True,
            token=HF_TOKEN
        )
        with download_lock:
            downloaded_components["tokenizer"] = tokenizer


    torch.cuda.empty_cache()


    with ThreadPoolExecutor(max_workers=2) as executor:
        model_future = executor.submit(download_base_model)
        tokenizer_future = executor.submit(download_tokenizer)


        model_future.result()
        tokenizer_future.result()

    model = downloaded_components["model"]
    tokenizer = downloaded_components["tokenizer"]


    torch.cuda.empty_cache()


    try:
        adapter_path = model_id
        print(f"Loading adapter from {adapter_path}")
        model.load_adapter(adapter_path, "default", token=HF_TOKEN)
        print("Adapter loaded successfully")
    except Exception as e:
        print(f"Error loading adapter: {e}")
        raise


    model.config.use_cache = True
    model.eval()


    torch.cuda.empty_cache()

    return model, tokenizer

def run_inference(model, tokenizer, tokenized_inputs, generation_config, batch_size=4):
    results = []

    for i in tqdm(range(0, len(tokenized_inputs), batch_size)):
        batch = tokenized_inputs[i:i+batch_size]

        prompts = [
            f"""<start_of_turn>system
簡潔に回答してください。装飾や特殊記号は使用しないでください。
<end_of_turn>
<start_of_turn>user
{item["input"]}
<end_of_turn>
<start_of_turn>model
""" for item in batch
        ]


        inputs = tokenizer(
            prompts,
            padding=True,
            truncation=True,
            return_tensors="pt"
        ).to(model.device)

        with torch.no_grad(), autocast(dtype=torch.bfloat16):
            outputs = model.generate(
                **inputs,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                **generation_config
            )

            for idx, output in enumerate(outputs):
                response = tokenizer.decode(output, skip_special_tokens=True)

                if 'model\n' in response:
                    response = response.split('model\n')[-1].strip()
                elif 'model' in response:
                    response = response.split('model')[-1].strip()


                response = post_process_output(response)

                results.append({
                    "task_id": batch[idx]["task_id"],
                    "input": batch[idx]["input"],
                    "output": response
                })


        del outputs, inputs
        torch.cuda.empty_cache()

    return results

def post_process_output(response):
    response = response.strip()
    symbols_to_replace = ['**', '`', '|', '```', '---', '===']
    for symbol in symbols_to_replace:
        response = response.replace(symbol, ' ')
    return ' '.join(response.split())

GENERATION_CONFIG = {
    "max_new_tokens": 512,
    "use_cache": True,
    "do_sample": False,
    "num_beams": 4,
    "repetition_penalty": 1.2,
    "length_penalty": 1.0,
    "early_stopping": False
}

def load_input_data(file_path):
    tokenized_inputs = []
    with open(file_path, "r") as f:
        for line in f:
            if line.strip():
                dt = json.loads(line)
                tokenized_inputs.append({
                    "task_id": dt["task_id"],
                    "input": dt["input"]
                })
    return tokenized_inputs

def save_results(results, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    jsonl_path = os.path.join(output_dir, "Output.jsonl")

    with open(jsonl_path, 'w', encoding='utf-8') as f:
        for item in results:
            json.dump(item, f, ensure_ascii=False)
            f.write('\n')

    print(f"Saved results to: {jsonl_path}")

def main():
    model, tokenizer = load_model_and_tokenizer()
    tokenized_inputs = load_input_data("/content/elyza-tasks-100-TV_0.jsonl")
    results = run_inference(model, tokenizer, tokenized_inputs, GENERATION_CONFIG)
    save_results(results, "output")

if __name__ == "__main__":
    main()