rinnaをCTranslate2から遊ぶメモ

conda create -n ct2test python=3.10
activate ct2test

pip install ctranslate2 
pip install sentencepiece transformers protobuf==3.20.0
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

import subprocess

base_model = "rinna/japanese-gpt-neox-3.6b-instruction-ppo"
output_dir = "rinnna-ppo-int8"
quantization_type = "int8"

command = f"ct2-transformers-converter --model {base_model} --quantization {quantization_type} --output_dir {output_dir}"
subprocess.run(command, shell=True)

import ctranslate2
import transformers
import time

base_model = "rinna/japanese-gpt-neox-3.6b-instruction-ppo"
quant_model = "rinnna-ppo-int8"

generator = ctranslate2.Generator(quant_model, device="cuda")
tokenizer = transformers.AutoTokenizer.from_pretrained(base_model, use_fast=False)

# プロンプトを作成する
def prompt(msg):
    p = [
        {"speaker": "ユーザー", "text": msg},
    ]
    p = [f"{uttr['speaker']}: {uttr['text']}" for uttr in p]
    p = "<NL>".join(p)
    p = p + "<NL>" + "システム: "
    return p

# 返信を作成する
def reply(msg):
    tokens = tokenizer.convert_ids_to_tokens(tokenizer.encode(prompt(msg),add_special_tokens=False,))

    results = generator.generate_batch(
        [tokens],
        max_length=256,
        sampling_topk=10,
        sampling_temperature=0.9,
        include_prompt_in_result=False,
    )

    text = tokenizer.decode(results[0].sequences_ids[0])
    print("A: " + text)
    return text

if __name__ == "__main__":
    while True:
        msg = input("Q: ")
        start_time = time.time()
        reply(msg)
        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"Elapsed time: {elapsed_time} seconds")