LLM-jp-4を動かしたりRAGを試したりした記録

お疲れ様です。

Gemma4と同時期くらいに日本語特化のLLMであるLLM-jp-4が出ていました。こちらも試してみたのでそれをまとめます。

LLM-jp-4について

国立情報学研究所のLLM研究のグループが開発した新たな国産LLMです。日本語性能でGPT-4oを上回る性能を出したとのこと。
ライセンスは「Apache-2.0」なので利用もしやすいですね。今後このモデルをベースに新たな日本語LLMが出てくるのも期待できそうです。

www.nii.ac.jp

実装

モデルはHuggingFaceで公開されているのみなので、今回はこちらを使用します。 Ollamaで使えれば楽なのですがこればかりは仕方ないですね…。

huggingface.co

ソースコード

ソースコードはいつもOllamaのモデルの検証に使用しているリポジトリを使用しました。こちらにHuggingFaceのモデル対応の処理を新たに作成しています。

github.com

とりあえず動かしてみる

まずはLLM-jp-4のモデルをそのまま動かしてみました。モデルページのUsageの部分のコードをベースにStreamlitのチャットUIで動くように改修しました。公式が公開しているcookbookにも同様にサンプルコードがあるようです。

素の状態の出力だと以下のように「<|channel|>」のような区切りトークンを含む出力になります。

output

区切りトークンを含む生成文をパースしてモデルの生成文のみを抜き出すようにそれ用の関数を作成して対応しました。

output_parsed

コード

def response_generator_huggingface_model() -> str:
    """HuggingFaceのモデルを使用
    """
    tokenizer = AutoTokenizer.from_pretrained(
        HF_MODEL,
        trust_remote_code=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        HF_MODEL,
        dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.eval()
    
    system_prompt = [{
        "role": "system",
        "content": "あなたはユーザの質問に答えるアシスタントです。回答は200文字程度で要点だけをまとめて簡潔に答えてください。"
    }]
    
    prompt: str = tokenizer.apply_chat_template(
        system_prompt + st.session_state.messages,
        tokenize=False,
        add_generation_prompt=True,
        reasoning_effort="low",  # {"low", "medium", "high"}
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        output_tensor = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
        )
    
    # 出力データを変換して返答文を取得
    generated_ids: list[int] = output_tensor[0][inputs["input_ids"].shape[1]:].tolist()
    response = tokenizer.decode(generated_ids)
    parsed_response = parse_chat_output(response)
    
    response_html = Markdown().convert(parsed_response["assistant"]["message"])
    
    return response_html

RAGを試してみる

次にいつものようにRAGも試してみました。

今回はHuggingFaceのモデルしかないので、最初はLangChainのHuggingFacePipelineを使って実装していました。ただこれがうまくいかず…。
以下にスクショを載せますが生成された返答文に与えたコンテキストがすべて残っておりめちゃくちゃ長くなっています。加えて先ほど素の状態で動かしたときのような区切りトークンになっておらず、出力のたびに形式が異なるので対応が難しい状態でした…。

rag_output_1

コード（改良前）

def response_generator_langchain_huggingface_rag() -> str:
    """langchain_huggingfaceを使用+RAG
    """
    tokenizer = AutoTokenizer.from_pretrained(
        HF_MODEL,
        # trust_remote_code is required to load custom tokenizer and reasoning parser.
        trust_remote_code=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        HF_MODEL,
        dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    pipe = pipeline(
        "text-generation", model=model, tokenizer=tokenizer, 
        max_new_tokens=256, do_sample=True, temperature=0.7, top_p=0.9
    )
    llm = HuggingFacePipeline(pipeline=pipe)
    
    # 直前のユーザの入力を取得
    user_input = st.session_state.messages[-1]["content"]
    
    # ベクトル化する準備
    model_kwargs = {
        "device": "cpu", # NOTE: モデルと合わせてVRAM容量を超えるのでCPUで実行
        "trust_remote_code": True
    }
    embedding = HuggingFaceEmbeddings(
        model_name="pfnet/plamo-embedding-1b",
        model_kwargs=model_kwargs
    )
    
    # DBを読み込んで知識データ取得
    vectorstore = Chroma(collection_name="elephants", 
                         persist_directory=DATABASE_DIR, 
                         embedding_function=embedding)
    docs = vectorstore.similarity_search(query=user_input, k=10)
    context = "\n".join([f"Content:\n{doc.page_content}" for doc in docs])
    
    messages = [
        ROLES[msg["role"]](content=msg["content"]) 
        for msg in st.session_state.messages[:-1]
    ] + [HumanMessage(content=RAG_PROMPT.format(question=user_input, context=context))]
    
    response = llm.invoke(messages)
    
    response_html = Markdown().convert(response)
    
    return response_html

tokenizerのchat_templateを使うと良いとの情報を得たので、LangChainを使わず素の状態のモデルを動かしたときのコードをベースにRAG用コードに改修して改良版を作成してみました。そうすることで、出力の形式が統一され生成文をパースできるようになりました。
以下のように返答文の部分だけうまく取り出せています。システムプロンプトも加えているので生成文の形式が少し変わっています。

rag_output_2

生成にかかる時間も大体1分ちょっとくらいです。前回の記事で比較したgpt-oss:20bやgemma4:e4bと一応変わらないくらいの処理時間ではありますね。

コード（改良後）

def response_generator_langchain_huggingface_rag() -> str:
    """langchain_huggingfaceを使用+RAG
    """
    tokenizer = AutoTokenizer.from_pretrained(
        HF_MODEL,
        # trust_remote_code is required to load custom tokenizer and reasoning parser.
        trust_remote_code=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        HF_MODEL,
        dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.eval()
    
    # 直前のユーザの入力を取得
    user_input = st.session_state.messages[-1]["content"]
    
    # ベクトル化する準備
    model_kwargs = {
        "device": "cpu", # NOTE: モデルと合わせてVRAM容量を超えるのでCPUで実行
        "trust_remote_code": True
    }
    embedding = HuggingFaceEmbeddings(
        model_name="pfnet/plamo-embedding-1b",
        model_kwargs=model_kwargs
    )
    
    # DBを読み込んで知識データ取得
    vectorstore = Chroma(collection_name="elephants", 
                         persist_directory=DATABASE_DIR, 
                         embedding_function=embedding)
    docs = vectorstore.similarity_search(query=user_input, k=10)
    context = "\n".join([f"Content:\n{doc.page_content}" for doc in docs])
    
    # システムプロンプトの用意
    system_prompt = [{
        "role": "system",
        "content": "あなたはユーザの質問に答えるアシスタントです。回答は最大500文字でまとめて簡潔に答えてください。"
    }]
    
    # 直前のユーザの入力を取得
    rag_input = [{
        "role": "user",
        "content": RAG_PROMPT.format(question=user_input, context=context)
    }]
    
    prompt: str = tokenizer.apply_chat_template(
        system_prompt + st.session_state.messages[:-1] + rag_input,
        tokenize=False,
        add_generation_prompt=True,
        reasoning_effort="low",  # {"low", "medium", "high"}
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        output_tensor = model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
        )
    
    # 出力データを変換して返答文を取得
    generated_ids: list[int] = output_tensor[0][inputs["input_ids"].shape[1]:].tolist()
    response = tokenizer.decode(generated_ids)
    parsed_response = parse_rag_output(response)
    
    response_html = Markdown().convert(parsed_response)
    
    return response_html