llm-jpをColabで試す

!pip install transformers accelerate sentencepiece --quiet

%time

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "llm-jp/llm-jp-13b-instruct-full-dolly-oasst-v1.0"

tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    device_map="cuda:0",
    torch_dtype=torch.float16,
).eval()

tokenizer.vocab_size

generation_config = {
    "max_new_tokens": 256,
    "do_sample": True,
    "temperature": 0.7,
    "top_p": 0.95,
}

text = "自然言語処理とは何か"
text = text + "### 回答："

with torch.no_grad():
    inputs = tokenizer(text, add_special_tokens=False, return_tensors="pt").to(model.device)

    outputs = model.generate(
        **inputs,
        **generation_config,
    )

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

def format_prompt(
    prompt: str,
    system_prompt: str = "" # 今回特に必要ない
) -> str:
    prompt = prompt + "### 回答： "
    prompt = system_prompt + prompt
    return prompt

format_prompt("１＋１は何？")

def ask(
    prompt: str,
    system_prompt: str | None = "",
    **kwargs
) -> str:
    generation_config = {
        "max_new_tokens": 128,
        "do_sample": True,
        "temperature": 0.7,
        "top_p": 0.95,
    }
    generation_config.update(kwargs)

    with torch.no_grad():
        prompt = format_prompt(prompt, system_prompt)
        inputs = tokenizer(
            prompt,
            add_special_tokens=False,
            return_tensors='pt',
        ).to(model.device)

        # .to(model.device)
        outputs = model.generate(
            **inputs,
            **generation_config,
        )

    output = tokenizer.decode(outputs[0])

    print(output)
    return output

ask("六本木周辺の観光スポットを教えてください。");

text = """
りんごが5つあります。そこから2つのりんごを取り除きました。残りのりんごの数は何個でしょう？
""".strip()

ask(text);

text = """バットとボールの両方を買うと1100円です。バットはボールよりも1000円高いです。ボールはいくらでしょう？""".strip()

ask(prompt=text);

text = """
引数kを取り、返り値としてフィボナッチ数列におけるk個目の値を返すPython関数を書いてください。
""".strip()

ask(prompt=text);

text = """
下記の英語を日本語に翻訳してください。
`There were 3 apples and 2 oranges. How many fruits were there in total?`
""".strip()

ask(prompt=text);

text = """
下記の文章を要約してください。

``

【速報】ロシア月探査機「ルナ 25」が月に衝突 「消滅した」ロスコスモス発表
月に向かっていたロシアの無人探査機「ルナ 25」が月に衝突したことが分かった。ロシアの国営宇宙企業ロスコスモスは先ほど、「月に衝突し、消滅した」と明らかにした。月面着陸前の軌道に移行中、制御不能となったという。
探査機は 21 日に月の南極付近に着陸予定だった。

``
""".strip()

ask(prompt=text);

text = """
あなたは友達ボットです。できるだけユーザーが親近感を感じやすいよう接してください。

ユーザー: 今日バイト、クビにされたわー。あー人生の意味って何なんだろうねぇー。
アシスタント:
""".strip()

ask(prompt=text);

text = """
### Question
There was a cookie on the table.
Tom entered the room.
The cookie disappeared.
What was likely to have happened?
""".strip()

ask(prompt=text);

text = """
### 質問
テーブルにクッキーがおいてありました。
太郎が部屋に入りました。
クッキーが消えました。
何が起きた可能性が高いですか？
""".strip()

ask(prompt=text);

text = """
必ず関西弁で答えてください。
たこ焼きのレシピを教えてください。
""".strip()

ask(prompt=text);