このコードは、gemma-2-9bに、 "elyza/ELYZA-tasks-100"のデータセットを3epochFine-Tuningして作成した、本adapter(marimo2000/gemma-2-9b-elyza-100_x3_lora)を使った推論コードである。ただし、以下で示すように推論前にRAGを行い、場合に応じてプロンプトの変更を行なっている。

背景

elyza/ELYZA-tasks-100のベンチマークの検証により、いくつかのタスクでZero Shot Chain Of Thought(以下ZeroShotCoT)、すなわち"Let's think step by step"を付加することで、スコアの上昇が見られたタスクが存在した。

そのため、elyza/ELYZA-tasks-100のInput部分を用いてRAGのコレクションを作成し、ZeroShotCoTが有効なタスクと類似したタスクが存在した場合、ZeroShotCoTのプロンプトを付加してスコアの向上を狙う戦略で進めた。

（NOTE：elyza/ELYZA-tasks-100のベンチマークでは、ZeroShotCoTを付加した場合、スコア上昇するタスクよりスコアが減少するタスクの方が多かった。）

残念ながら、ELYZA-tasks-100-TVの最終スコアは、ZeroShotCoTを使ったもののほうが、0.07ポイントほど下がってしまったが、RAGのチューニングを行えば、

ZeroShotCoTの有効性をより発揮できる可能性があり、試行錯誤の結果として本コードを最終版として提出することに決めた。

! python --version

! nvidia-smi

python

# 必要なライブラリをインストール
! pip install -U -q unsloth
! pip install -U -q torch
! pip install -U -q peft
! pip install chromadb -U -q
! pip install openpyxl -U -q

python

# 必要なライブラリを読み込み
from unsloth import FastLanguageModel
from peft import PeftModel
import json
from tqdm import tqdm
import re
import pandas as pd
from datasets import load_dataset

import os
from dotenv import load_dotenv
load_dotenv()

import chromadb

python

from enum import Enum

class TASK(Enum):
    ELYZA_100_TV = "elyza100tv"
    ELYZA_100 = "elyza100"


class Prompt(Enum):
    Defualt = "default"
    ZeroShotCoT = "ZS-CoT"

RAGの準備

python

collection_name = 'elyza100'

persist_path = './outputs/rag/'

chroma_client = chromadb.PersistentClient(path=persist_path)

RAGのCollectionの作成

python

try:
    # 再現性のためにコレクションを削除して再作成
    chroma_client.delete_collection(name=collection_name) 
except:
    print('collection not found')
collection = chroma_client.get_or_create_collection(collection_name)

python

def load_datasets(inference_task: TASK)->pd.DataFrame:
    """
    inference_taskに応じたデータセットを読み込む関数
    """
    # タスクとなるデータの読み込み。
    # 事前にデータをアップロードしてください。
    if inference_task == TASK.ELYZA_100_TV:
        datasets = []
        with open("./elyza-tasks-100-TV_0.jsonl", "r") as f:
            item = ""
            for line in f:
                line = line.strip()
                item += line
                if item.endswith("}"):
                    datasets.append(json.loads(item))
                    item = ""
        datasets = pd.DataFrame(datasets)
    # Elyza100のタスクを解く場合
    elif inference_task == TASK.ELYZA_100:
        datasets_raw = load_dataset("elyza/ELYZA-tasks-100", revision="1.0.0")

        datasets = datasets_raw["test"].to_pandas()

        datasets["task_id"] = datasets.index.astype(str)
    else:
        raise ValueError("Invalid task")

    return datasets

python

# elyza/ELYZA-tasks-100 データセットの読み込み

elyza100_data = load_datasets(TASK.ELYZA_100)

python

collection.add(
documents=elyza100_data["input"].to_list(),
ids= elyza100_data.index.map(lambda x: str(x)).to_list(),)

RAGのQuery用のヘルパー関数

python

def get_answer_from_result(idx, result):
    """
    RAGの検索結果からidxに応じた回答を取得する
    """
    return {
        "id": result["ids"][0][idx],
        "distance": round(result["distances"][0][idx], 2),
        "document": result["documents"][0][idx][:100]
            }
def convert_result_to_df(result):
    """
    RAGの検索結果をDataFrameに変換する
    """
    df = []
    for idx in range(len(result["ids"][0])):
        df.append(get_answer_from_result(idx, result))
    return pd.DataFrame(df)

input = "仕事の熱意を取り戻すためのアイデアを5つ挙げてください。"
result=collection.query(
        query_texts=[input],
        n_results=10,
    )

df = convert_result_to_df(result)
df

Fine-TuningしたLLMの準備

NOTE: ZeroShotCoTが有効なElyza100のtask_idはcsvにまとめて本README.mdのディレクトリに保存している（gemma-2-9b-elyza-100_x3_lora_diff_high_v1.csv)

以下のCOT_COEFF_TASK_ID_FILE_PATHでファイルパスを宣言する

python

# ベースとなるモデルと学習したLoRAのアダプタ（Hugging FaceのIDを指定）。
model_id = "google/gemma-2-9b"
adapter_id = "marimo2000/gemma-2-9b-elyza-100_x3_lora"

model_local_path = f'content/{os.path.basename(model_id)}'

HF_TOKEN = os.getenv("HF_TOKEN")

COT_COEFF_TASK_ID_FILE_PATH = "outputs/elyza100/score/gemma-2-9b-elyza-100_x3_lora_diff_high_v1.csv"

VERSION = "FINAL"

NOTE: モデルをLocalにダウンロードして、そこから読み込みFastLanguageModelを作成する。

python

# すでにモデルが存在する場合は終了
if os.path.exists(model_local_path):
    print("Model already exists")
else:
    os.makedirs(model_local_path, exist_ok=True)
    !huggingface-cli login --token {HF_TOKEN}
    !huggingface-cli download {model_name} --local-dir={model_local_path}

    !ls -lh {model_local_path}

python

# unslothのFastLanguageModelでLocalのモデルを読み込む
dtype = None # Noneにしておけば自動で設定
load_in_4bit = True # 今回は13Bモデルを扱うためTrue

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_local_path,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
    trust_remote_code=True,
)

python

# 元のモデルにLoRAのアダプタを統合。
model = PeftModel.from_pretrained(model, adapter_id, token = HF_TOKEN)

python

# モデルを用いてタスクの推論。

# 推論するためにモデルのモードを変更
FastLanguageModel.for_inference(model)

RAG及びLLMによる推論用の関数の定義

python

def load_cot_efficiency_task_ids(file_path: str)->list:
    """
    CoTの効率化が必要なタスクidを読み込む関数
    """
    df_cot_efficiency_task_ids = pd.read_csv(file_path)
    cot_efficiency_task_ids = df_cot_efficiency_task_ids["task_id"].tolist()
    cot_efficiency_task_ids.sort()
    print(f"CoTの効率化が必要なタスクid: {cot_efficiency_task_ids}")
    return cot_efficiency_task_ids

def save_results(results, file_path):
    """
    結果を保存する関数
    """
    with open(file_path, "w", encoding='utf-8') as f:
        for result in results:
            json.dump(result, f, ensure_ascii=False)
            f.write('\n')

    print(f"結果を{file_path}に保存しました。")

    

def do_task(model, adapter_id, tokenizer, collection, inference_task: TASK):
    """
    タスクを解く関数
    """

    print(f"推論するタスクは{inference_task.value}です。")

    # ここではadapter_idを元にファイル名を決定しているが、ファイル名は任意で問題なし。
    output_dir = f"outputs/{inference_task.value}"
    json_file_id = re.sub(".*/", "", adapter_id)
    results_file_path = f"{output_dir}/{json_file_id}_with_rag_output_v{VERSION}.jsonl"

    print(f"結果は{results_file_path}に保存されます。")

    # データセットの読み込み
    datasets = load_datasets(inference_task)

    # CoTの効率化が必要なタスクのidを読み込む
    cot_efficiency_task_ids = load_cot_efficiency_task_ids(file_path=COT_COEFF_TASK_ID_FILE_PATH)

    results = []
    for _, dt in tqdm(datasets.iterrows(), total=len(datasets)):
        input = dt["input"]
        
        # CoTの効率化が必要なタスクの場合はプロンプトを変更
        if is_cot_valid_question(input, collection, cot_efficiency_task_ids):
            prompt = get_prompt(Prompt.ZeroShotCoT, input)
        else:
            prompt = get_prompt(Prompt.Defualt, input)

        # LLMによる生成
        inputs = tokenizer([prompt], return_tensors = "pt").to(model.device)
        outputs = model.generate(**inputs, max_new_tokens = 512, use_cache = True, do_sample=False, repetition_penalty=1.2)
        prediction = tokenizer.decode(outputs[0], skip_special_tokens=True).split('### 回答\n')[-1]
        
        # 結果をListに追加
        results.append({"task_id": dt["task_id"], "input": input, "output": prediction})

    # 結果を保存
    save_results(results, results_file_path)

def get_prompt(prompt: Prompt, input: str):
    """
    prompのenumに応じたプロンプトを返す。
    """
    if prompt == Prompt.Defualt:
        return f"""### 指示\n{input}\n### 回答\n"""
    elif prompt == Prompt.ZeroShotCoT:
        return f"""### 指示\n{input}\nLet's think step by step\n### 回答\n"""
    else:
        raise ValueError("Invalid prompt")

def is_cot_valid_question(input, collection, cot_efficiency_task_ids: list):
    """
    CoTの効率化が必要なタスクかどうかを判定する。
    """
    results =  collection.query(
        query_texts=[input],
        n_results=1,
    )

    #top1の中にcot_efficiency_task_idsが含まれているか

    df = convert_result_to_df(results)
    for idx, row in df.iterrows():
        if int(row["id"]) in cot_efficiency_task_ids:
            return True
    
    return False

cot_efficiency_task_ids = load_cot_efficiency_task_ids(file_path=COT_COEFF_TASK_ID_FILE_PATH)

do_task(model, adapter_id, tokenizer, collection, TASK.ELYZA_100_TV)

marimo2000/llm-jp-3-13b-elyza-100-it_lora

Uploaded model

HOW TO USE THIS MODEL

推論コード

概要

背景

RAGの準備

RAGのCollectionの作成

RAGのQuery用のヘルパー関数

Fine-TuningしたLLMの準備

NOTE: モデルをLocalにダウンロードして、そこから読み込みFastLanguageModelを作成する。

RAG及びLLMによる推論用の関数の定義

marimo2000/
llm-jp-3-13b-elyza-100-it_lora