AWS LambdaでGoogleのAIモデル「Gemma」を動かす

2024年2月21日、Googleから大規模言語モデル(LLM)である「Gemma(ジェマ)」が公開されました！
これまでGoogleのモデルはクローズドなものでしたが、このGemmaはオープンソース型！ということで、本記事ではAWS Lambdaを使用してGemmaを実行させてみます。

## はじめに
Gemmaの公開から1週間ちょっと経っていますが、この短い間に既に色々と検証されている方や、Gemmaを利用するためのライブラリ等の対応など多くのことがネット上に見受けられます。
特に、GemmaをGGUF形式に変換するllama.cppはGemmaのリリースから3日と経たずGemmaに対応されており、最近のLLMに関する開発スピードの早さには驚かされます。
ちなみにこの辺りの対応の流れは以下のURLから確認いただけるので、ご興味があれば覗いてみてください。

https://github.com/ggerganov/llama.cpp/pull/5631

## 手順
手順としては大きく以下のようになります。
1. Gemma(GGUF形式のモデル)のダウンロード
1. Lambdaのコード類の準備
1. Lambda作成＆テスト

それでは順番にみていきます。

## 1. Gemma(GGUF形式のモデル)のダウンロード
先ずはGemmaのモデルをダウンロードします。

## ライセンスへの同意
Gemmaは先に記述したとおりオープンソース型ではありますが、使用するにはライセンスを確認して同意する必要があります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2741318/60eadad5-c7a7-baed-9dd2-2b42fa5c9d44.png)

Hugging Faceにログインし、フォームに氏名等を入力して処理を進めることでモデルへのアクセスが出来るようになります。

https://huggingface.co/google/gemma-7b-it/tree/main

## モデルダウンロード
それでは、Lambdaで利用するGemmaのモデルをダウンロードします。

今回はLambdaで動作させることを目指していることから、サイズ等についてやや厳しめの制限があります。
そのため必然的に軽量なモデルが必要になるのですが、残念ながらGoogleがデフォルトで提供しているモデルだと何れもサイズが大きくそのままでは利用が難しいです。
そのため本記事では量子化して軽量化されたGGUF形式のモデルを使用します。

既にいくつもGGUF形式のモデルが公開されていますが、今回は以下のURLから「gemma-7b-it-Q4_K_M.gguf」をダウンロードして使用させてもらいます。

https://huggingface.co/rahuldshetty/gemma-7b-it-gguf-quantized/tree/main

ダウンロードしたモデルは任意のディレクトリに格納してください。
後ほど同ディレクトリにはLambdaのためのソースコード類を格納します。

## 2. Lambdaのコード類の準備

本記事では最終的に次のようなファイル構成になります。

```shell
lambda-gemma
├── app.py
├── Dockerfile
├── gemma-7b-it-Q4_K_M.gguf
├── install.sh
└── uninstall.txt
```
gemma-7b-it-Q4_K_M.ggufファイルは先の手順でダウンロードしたモデルです。
それ以外の4つのファイルについては以下に記載します。

app.pyはLambdaのコードになります。
```python:app.py
# -- coding: utf-8 --`
import gc
from llama_cpp import Llama

model = Llama(model_path="/var/model/gemma-7b-it-Q4_K_M.gguf")

DEFAULT_PROMPT = "Any recommendations for Japanese cuisine?"
DEFAULT_MAX_TOKENS = 512

def handler(event, context):
    global model

    prompt = event.setdefault('prompt', DEFAULT_PROMPT)
    mt = event.setdefault('max_tokens', DEFAULT_MAX_TOKENS)

    output = model(prompt, max_tokens=mt)

    gc.collect()
    return  {"statusCode": 200, "body": {"prompt": prompt, "output": output}}
```

コンテナ形式のLambdaを作成するためのDockerfileになります。
```Dockerfile:Dockerfile
FROM python:3.12 as production
RUN python -m pip install --upgrade pip && pip install awslambdaric llama-cpp-python==0.2.53 && pip cache purge
COPY gemma-7b-it-Q4_K_M.gguf /var/model/
COPY app.py ${LAMBDA_TASK_ROOT}/

ENTRYPOINT [ "/usr/local/bin/python", "-m", "awslambdaric" ]
CMD [ "app.handler" ]
```

AWS ECRの作成からDockerfileのビルド、ビルドしたイメージのプッシュなどを実行するためのスクリプトになります。
```shell:install.sh
#!/bin/bash
set -euo pipefail

echo -n "(Create New) Input AWS Lambda Function Name [ex. myLLMFunction]: "
read -r LAMBDANAME

REGION=$(aws configure get region)
ACCOUNTID=$(aws sts get-caller-identity --output text --query Account)
RAND=$(date +%Y%m%d%H%M%S%3N | shasum -a 512 | base64 | fold -w 16 | head -n 1 | tr 'A-Z' 'a-z')

# Lambda function Name
if [[ -z ${LAMBDANAME} ]]; then
  LAMBDANAME="myLLMFunction"
fi
LAMBDANAME="$LAMBDANAME-$RAND"
# Role Name
ROLENAME="$LAMBDANAME-role"
# ECR Repository - lambda container image stored
REPOSITORYNAME=$(echo "$LAMBDANAME-repo" | tr 'A-Z' 'a-z')

print_message() {
  echo "[$(date +'%Y-%m-%d %H:%M:%S')] $1"
}

create_ecr_repository() {
  local result
  result=$(aws ecr create-repository --repository-name "${REPOSITORYNAME}")
  if grep -q 'repositoryArn' <<< "${result}"; then
    print_message "Success!!: Create Repository"
  else
    print_message "Error occurred: ${REPOSITORYNAME}"
    exit 1
  fi
}

build_and_push_container() {
  local repo_uri="${ACCOUNTID}.dkr.ecr.${REGION}.amazonaws.com/${REPOSITORYNAME}"
  aws ecr get-login-password --region "${REGION}" | docker login --username AWS --password-stdin "${ACCOUNTID}.dkr.ecr.${REGION}.amazonaws.com"
  docker build -t "${REPOSITORYNAME}" .
  docker tag "${REPOSITORYNAME}:latest" "${repo_uri}:latest"
  docker push "${repo_uri}:latest"
}

create_role() {
  local result
  result=$(aws iam create-role --role-name ${ROLENAME} --assume-role-policy-document '{"Version": "2012-10-17","Statement": [{ "Effect": "Allow", "Principal": {"Service": "lambda.amazonaws.com"}, "Action": "sts:AssumeRole"}]}')
  aws iam attach-role-policy --role-name ${ROLENAME} --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
}

create_lambda_function() {
  DIGEST=$(aws ecr list-images --repository-name ${REPOSITORYNAME} --out text --query 'imageIds[?imageTag==`latest`].imageDigest')
  aws lambda create-function --function-name ${LAMBDANAME} --out text \
      --package-type Image --code ImageUri=${ACCOUNTID}.dkr.ecr.${REGION}.amazonaws.com/${REPOSITORYNAME}@${DIGEST} \
      --memory-size 10240 --timeout 900 \
      --role arn:aws:iam::${ACCOUNTID}:role/${ROLENAME}
}

write_uninstaller() {
  cat ./uninstall.txt | sed -e "s/BEFORELAMBDANAME/${LAMBDANAME}/" \
    -e "s/BEFOREROLENAME/${ROLENAME}/" \
    -e "s/BEFOREREPOSITORYNAME/${REPOSITORYNAME}/" \
    > "./uninstall-${LAMBDANAME}.sh"
}

print_message "Starting script..."

print_message "(1/5) Create AWS ECR Repository"
create_ecr_repository

print_message "(2/5) Create Role"
create_role

print_message "(3/5) Build and Push Container"
build_and_push_container

print_message "(4/5) Create Lambda"
create_lambda_function

print_message "(5/5) Writing Uninstaller"
write_uninstaller

print_message "******* Complete!! *******"
print_message "The following resources were created."
print_message "- Lambda function: ${LAMBDANAME}"
print_message "- Role: ${ROLENAME}"
print_message "- ECR Repository: ${REPOSITORYNAME}"
```

先のスクリプトでAWS上に作成した各種リソースを削除するためのスクリプトを作成するためのファイルになります。
```shell:uninstall.txt
#!/usr/bin/bash
aws iam detach-role-policy --role-name BEFOREROLENAME --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
aws iam delete-role --role-name BEFOREROLENAME
aws lambda delete-function --output text --function-name BEFORELAMBDANAME
aws ecr delete-repository --output text --repository-name BEFOREREPOSITORYNAME --force
```

## 3. Lambda作成＆テスト

次にLambda関数の作成等を行いますが、ここでは先ほど用意した「install.sh」を使用します。
ターミナルからこのスクリプトを実行していただきますが、事前にAWS CLIの設定をお願いします。

以下のようにinstall.shを実行すると、Lambdaの関数名を聞かれるため任意の関数名を入力してください。
```shell:CMD
 $ bash ./install.sh
```
スクリプトではコンテナイメージのビルドやプッシュなどを行うため完了には時間がかかります。
また、スクリプト実行完了後もLambda関数が実行可能になるまでにも多少時間がかかります。

### 動作確認
それではLambda関数のテストです。
プロンプトの指定は、イベントJSONで次のように設定します。

```json:イベントJSON
{
  "prompt": "Any recommendations for Japanese cuisine?"
}
 ```
 
ここではテスト用のプロンプトは英語で入力しています。
()内の日本語はこちらで別途訳したものになります。
> Any recommendations for Japanese cuisine?
> (おすすめの日本料理はありますか？)

このプロンプトに対する回答は以下のとおりです。

> Whether you're a novice or an aficionado, I have some suggestions to elevate your experience: 
> (初心者の方にも愛好家の方にも、あなたの経験をより良いものにするための提案があります：)
> **Novice:**
> (**初心者：**)
> - **Donburi (Rice Bowl)** at the supermarket with tempura vegetables and marinated pork. This is accessible & affordable way introduce yourself! (Example : Poke Sushi)
> (- 天ぷら野菜と豚肉のマリネがのったスーパーのどんぶり。身近で手頃な自己紹介！(例：ポケ寿司）)

何度か試すと普通に寿司やテンプラを回答してくれることもありますが、このような変化球の回答もままあります。
Lambda関数実行から回答までに要した時間はだいたい数分程度でした。

ここまでで動作確認できたので、以上です。