Google Colab で CodeGemma を試す
「Google Colab」で「CodeGemma」を試したので、まとめました。
1. CodeGemma
「CodeGemma」は、コードタスク用のモデルです。次の3種類のモデルが提供されています。
2. コード補完
Colabでのコード補完の手順は、次のとおりです。
(1) パッケージのインストール。
# パッケージのインストール
!pip install transformers accelerate
(2) 「HuggingFace」からAPIキーを取得し、Colabのシークレットマネージャーの「HF_TOKEN」に登録。
(3) トークナイザーとモデルの準備。
今回は、「google/codegemma-2b」を使います。
from transformers import GemmaTokenizer, AutoModelForCausalLM
# トークナイザーとモデルの準備
tokenizer = GemmaTokenizer.from_pretrained(
"google/codegemma-2b"
)
model = AutoModelForCausalLM.from_pretrained(
"google/codegemma-2b",
device_map="auto",
torch_dtype="auto",
)
(4) 推論の実行。
# プロンプトの準備
prompt = '''\
<|fim_prefix|>import datetime
def calculate_age(birth_year):
"""Calculates a person's age based on their birth year."""
current_year = datetime.date.today().year
<|fim_suffix|>
return age<|fim_middle|>\
'''
# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
prompt_len = inputs["input_ids"].shape[-1]
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0][prompt_len:]))
age = current_year - birth_year<|file_separator|><eos>
プロンプトで使用するスペシャルトークンは、次のとおりです。
3. コード生成
(1) 推論の実行。
# プロンプトの準備
# n番目のフィボナッチ数を計算するPython関数を書いてください。
prompt = "Write me a Python function to calculate the nth fibonacci number."
# 推論の実行
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
prompt_len = inputs["input_ids"].shape[-1]
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0][prompt_len:]))
py
<|fim_prefix|><|fim_suffix|><|fim_middle|>def fibonacci(n):
▁▁▁▁if n == 0:
▁▁▁▁▁▁▁▁return 0
▁▁▁▁elif n == 1:
▁▁▁▁▁▁▁▁return 1
▁▁▁▁else:
▁▁▁▁▁▁▁▁return fibonacci(n-1) + fibonacci(n-2)
▁▁▁▁
n = int(input("Enter the nth term: "))
print("The nth term of the Fibonacci series is:", fibonacci(n))
<|file_separator|><eos>
この記事が気に入ったらサポートをしてみませんか?