استنتاج با جما با استفاده از JAX و Flax

مشاهده در ai.google.dev در Google Colab اجرا شود در Vertex AI باز کنید مشاهده منبع در GitHub

بررسی اجمالی

Gemma خانواده ای از مدل های سبک و پیشرفته زبان باز بزرگ است که بر اساس تحقیقات و فناوری Google DeepMind Gemini است. این آموزش نشان می دهد که چگونه می توان نمونه برداری/استنتاج اولیه را با مدل Gemma 2B Instruct با استفاده از کتابخانه gemma Google DeepMind که با JAX (یک کتابخانه محاسباتی عددی با کارایی بالا)، Flax (کتابخانه شبکه عصبی مبتنی بر JAX)، Orbax (یک) انجام داد. کتابخانه مبتنی بر JAX برای ابزارهای آموزشی مانند checkpointing) و SentencePiece (یک کتابخانه توکنایزر/دتوکنیزر). اگرچه از Flax مستقیماً در این نوت بوک استفاده نمی شود، از Flax برای ایجاد Gemma استفاده شده است.

این نوت بوک می تواند در Google Colab با GPU رایگان T4 اجرا شود (به Edit > تنظیمات نوت بوک > زیر شتاب دهنده سخت افزار T4 GPU را انتخاب کنید).

برپایی

1. دسترسی Kaggle را برای Gemma تنظیم کنید

برای تکمیل این آموزش، ابتدا باید دستورالعمل های راه اندازی را در Gemma setup دنبال کنید، که به شما نشان می دهد چگونه کارهای زیر را انجام دهید:

  • در kaggle.com به Gemma دسترسی پیدا کنید.
  • یک زمان اجرا Colab با منابع کافی برای اجرای مدل Gemma انتخاب کنید.
  • نام کاربری و کلید API Kaggle را ایجاد و پیکربندی کنید.

پس از تکمیل تنظیمات Gemma، به بخش بعدی بروید، جایی که متغیرهای محیطی را برای محیط Colab خود تنظیم خواهید کرد.

2. متغیرهای محیط را تنظیم کنید

متغیرهای محیطی را برای KAGGLE_USERNAME و KAGGLE_KEY تنظیم کنید. وقتی از شما خواسته شد "دسترسی بدهید؟" پیام‌ها، با ارائه دسترسی مخفی موافقت کنید.

import os
from google.colab import userdata # `userdata` is a Colab API.

os.environ["KAGGLE_USERNAME"] = userdata.get('KAGGLE_USERNAME')
os.environ["KAGGLE_KEY"] = userdata.get('KAGGLE_KEY')

3. کتابخانه gemma نصب کنید

این نوت بوک بر روی استفاده از یک GPU رایگان Colab تمرکز دارد. برای فعال کردن شتاب سخت افزاری، روی ویرایش > تنظیمات نوت بوک > انتخاب T4 GPU > ذخیره کلیک کنید.

در مرحله بعد، باید کتابخانه Google DeepMind gemma را از github.com/google-deepmind/gemma نصب کنید. اگر خطای «تحلیل کننده وابستگی پیپ» دریافت کردید، معمولاً می توانید آن را نادیده بگیرید.

pip install -q git+https://github.com/google-deepmind/gemma.git

مدل جما را بارگیری و آماده کنید

  1. مدل Gemma را با kagglehub.model_download بارگیری کنید که سه آرگومان می گیرد:
  • handle : مدل دسته از Kaggle
  • path : (رشته اختیاری) مسیر محلی
  • force_download : (بولی اختیاری) بارگیری مجدد مدل را مجبور می کند
GEMMA_VARIANT = '2b-it' # @param ['2b', '2b-it'] {type:"string"}

import kagglehub

GEMMA_PATH = kagglehub.model_download(f'google/gemma/flax/{GEMMA_VARIANT}')

Downloading from https://www.kaggle.com/api/v1/models/google/gemma/flax/2b-it/2/download...
100%|██████████| 3.67G/3.67G [00:35<00:00, 110MB/s]
Extracting model files...

print('GEMMA_PATH:', GEMMA_PATH)

GEMMA_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2
  1. محل وزن های مدل و توکنایزر را بررسی کنید، سپس متغیرهای مسیر را تنظیم کنید. دایرکتوری توکنایزر در دایرکتوری اصلی جایی که مدل را دانلود کرده‌اید قرار می‌گیرد، در حالی که وزن‌های مدل در یک دایرکتوری فرعی قرار دارند. مثلا:
  • فایل tokenizer.model در /LOCAL/PATH/TO/gemma/flax/2b-it/2 خواهد بود.
  • نقطه بازرسی مدل در /LOCAL/PATH/TO/gemma/flax/2b-it/2/2b-it خواهد بود.
CKPT_PATH = os.path.join(GEMMA_PATH, GEMMA_VARIANT)
TOKENIZER_PATH = os.path.join(GEMMA_PATH, 'tokenizer.model')
print('CKPT_PATH:', CKPT_PATH)
print('TOKENIZER_PATH:', TOKENIZER_PATH)

CKPT_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/2b-it
TOKENIZER_PATH: /root/.cache/kagglehub/models/google/gemma/flax/2b-it/2/tokenizer.model

انجام نمونه گیری/استنتاج

  1. با متد gemma.params.load_and_format_params ، چک پوینت مدل Gemma را بارگیری و فرمت کنید:
from gemma import params as params_lib

params = params_lib.load_and_format_params(CKPT_PATH)
  1. توکنایزر Gemma را که با استفاده از sentencepiece.SentencePieceProcessor ساخته شده است بارگیری کنید.SentencePieceProcessor:
import sentencepiece as spm

vocab = spm.SentencePieceProcessor()
vocab.Load(TOKENIZER_PATH)

True
  1. برای بارگیری خودکار پیکربندی صحیح از نقطه بازرسی مدل Gemma، از gemma.transformer.TransformerConfig استفاده کنید. آرگومان cache_size تعداد مراحل زمانی در حافظه پنهان Gemma Transformer است. پس از آن، مدل Gemma را به عنوان transformer با gemma.transformer.Transformer (که از flax.linen.Module به ارث می برد) نمونه سازی کنید.
from gemma import transformer as transformer_lib

transformer_config = transformer_lib.TransformerConfig.from_params(
    params=params,
    cache_size=1024
)

transformer = transformer_lib.Transformer(transformer_config)
  1. یک sampler با gemma.sampler.Sampler در بالای ایست بازرسی/وزن‌های مدل Gemma و توکنایزر ایجاد کنید:
from gemma import sampler as sampler_lib

sampler = sampler_lib.Sampler(
    transformer=transformer,
    vocab=vocab,
    params=params['transformer'],
)
  1. یک دستور در input_batch بنویسید و استنتاج را انجام دهید. می‌توانید total_generation_steps را تغییر دهید (تعداد مراحل انجام شده هنگام ایجاد یک پاسخ - این مثال 100 برای حفظ حافظه میزبان استفاده می‌کند).
prompt = [
    "\n# What is the meaning of life?",
]

reply = sampler(input_strings=prompt,
                total_generation_steps=100,
                )

for input_string, out_string in zip(prompt, reply.text):
    print(f"Prompt:\n{input_string}\nOutput:\n{out_string}")

Prompt:

# What is the meaning of life?
Output:


The question of what the meaning of life is one that has occupied the minds of philosophers, theologians, and individuals for centuries. There is no single, universally accepted answer, but there are many different perspectives on this complex and multifaceted question.

**Some common perspectives on the meaning of life include:**

* **Biological perspective:** From a biological standpoint, the meaning of life is to survive and reproduce.
* **Existential perspective:** Existentialists believe that life is not inherently meaningful and that
  1. (اختیاری) اگر نوت بوک را کامل کرده اید و می خواهید دستور دیگری را امتحان کنید، این سلول را اجرا کنید تا حافظه آزاد شود. پس از آن، می‌توانید sampler دوباره در مرحله 3 نمونه‌سازی کنید و درخواست را در مرحله 4 سفارشی کرده و اجرا کنید.
del sampler

بیشتر بدانید