vLLMでlightblue/qarasu-14B-chat-plus-unleashedをGoogleColobで動かす

AI

はじめに

年末にして新しいモデルが出たので、触っていきます

今回はvLLMで動かすサンプルコードがあったのでこちらで動かしていきます

環境

準備

vLLMで動かすので、以下を入れます

!pip install vllm
!pip install tiktoken

推論

モデルのロード

vLLMを使ってfloat16でモデルをロードします

from vllm import LLM, SamplingParams

sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
llm = LLM(model="lightblue/qarasu-14B-chat-plus-unleashed", trust_remote_code=True,dtype="float16")

サンプルプロンプト

推論

messages = [{"role": "system", "content": "あなたはAIアシスタントです。"}]
messages.append({"role": "user", "content": "イギリスの首相は誰ですか？"})
prompt = llm.llm_engine.tokenizer.apply_chat_template(conversation=messages, add_generation_prompt=True, tokenize=False)
prompts = [prompt]

outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

結果

'<|im_start|>system\nあなたはAIアシスタントです。<|im_end|>\n<|im_start|>user\nイギリスの首相は誰ですか？<|im_end|>\n<|im_start|>assistant\n', Generated text: '現在のイギリスの首相はリサ・ナウロンです。<|im_end|>\n'

まどマギプロンプト

推論

messages = [{"role": "system", "content": "あなたはアニメ評論家です"}]
messages.append({"role": "user", "content": "まどマギで一番可愛いキャラはなんですか？"})
prompt = llm.llm_engine.tokenizer.apply_chat_template(conversation=messages, add_generation_prompt=True, tokenize=False)
prompts = [prompt]

outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

結果

<|im_start|>system\nあなたはアニメ評論家です<|im_end|>\n<|im_start|>user\nまどマギで一番可愛いキャラはなんですか？<|im_end|>\n<|im_start|>assistant\n', Generated text: 'この文章は、2023年3月に書かれたものであり、当時の情報に基づいています。そのため、現在の状況や詳細については、正確さを保証するものではありません。\n\n2023年3月時点で、最も可愛いキャラクターは「アイドルマスター シンデレラガールズ」の「アイドルマスター シンデレラガールズ」の「アイドルマスター シン'

使用リソース