from huggingface_hub import snapshot_download  
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer, BitsAndBytesConfig
import torch 

def generate_response(model, tokenizer, prompt, device):
    """
    指定されたプロンプトに対して、モデルを用いてテキスト生成を実行し、
    生成されたテキスト応答を返却する関数です。

    Args:
        model: 使用する因果言語モデル (AutoModelForCausalLM)
        tokenizer: モデルに対応するトークナイザー (AutoTokenizer)
        prompt: テキスト生成の入力として与えるプロンプト文字列
        device: テンソルを転送するデバイス（例: "cuda" や "cpu"）

    Returns:
        生成されたテキスト応答（デコード済みの文字列）
    """
    inputs = tokenizer(prompt, return_tensors="pt") # テンソル変換
    inputs = {key: tensor.to(device) for key, tensor in inputs.items()} # デバイスに転送
    output = model.generate(**inputs, streamer=streamer) # TextStreamerを用いてテキスト生成を実行
    return tokenizer.decode(output[0], skip_special_tokens=True, max_new_tokens=1024) # 生成されたトークンをデコード

device = "cuda" # 使用するデバイスを "cuda" （GPU）に設定
prompt = "AIについて簡潔に教えてください。" # テキスト生成用の入力プロンプトを定義

# 量子化の設定: 8-bitロードを行い、4-bit計算時にtorch.float16を使用する設定を適用
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,                # モデルを4bitでロードする設定
    bnb_4bit_use_double_quant=True,   # 二重量子化の使用を指定
    bnb_4bit_quant_type="nf4",        # 量子化タイプを「nf4」に設定
    bnb_4bit_compute_dtype=torch.bfloat16,  # 計算時のデータ型をbfloat16に設定
)

# モデルをダウンロード
model_name = snapshot_download(
    repo_id="ArekuNoimar/llama-2-7b-hf-instruct-chatbot-arena-ja-v3" # モデルのリポジトリID
)

# 量子化設定およびデバイスマッピングを適用
model = AutoModelForCausalLM.from_pretrained(
    model_name,                      # ダウンロードしたモデルのパス
    device_map=device,               # テンソルを動作させるデバイス（"cuda"）を指定
    quantization_config=bnb_config,  # 設定した量子化設定を適用
    torch_dtype=torch.bfloat16       # モデルの計算に使用するデータ型をbfloat16に指定
)

tokenizer = AutoTokenizer.from_pretrained(model_name) # トークナイザーの読み込み
streamer = TextStreamer(tokenizer, skip_prompt=False, skip_special_tokens=False) # TextStreamerを初期化
generate_response(model, tokenizer, prompt, device) # 関数を呼び出してテキストを生成

yasutoshi-lab/llama-2-7b-hf-instruct-chatbot-arena-ja-v3

llama-2-7b-hf-instruct-chatbot-arena-ja-v3

モデルについて

注意点

使い方

yasutoshi-lab/
llama-2-7b-hf-instruct-chatbot-arena-ja-v3