প্রজন্ম এবং টিউনিং উদাহরণের জন্য জেমমা কুকবুক সংগ্রহস্থল দেখুন! আরও জানুন

এই পৃষ্ঠাটি Cloud Translation API অনুবাদ করেছে।

JAX এবং Flax ব্যবহার করে RecurrentGemma ফাইন-টিউনিং

ai.google.dev-এ দেখুন

Google Colab-এ চালান

Vertex AI-তে খুলুন

GitHub-এ উৎস দেখুন

এই টিউটোরিয়ালটি Google DeepMind এর recurrentgemma লাইব্রেরি , JAX (একটি উচ্চ-পারফরম্যান্স সংখ্যাসূচক কম্পিউটিং লাইব্রেরি), Flax (JAX-ভিত্তিক নিউরাল নেটওয়ার্ক লাইব্রেরি), Chex (একটি উচ্চ-পারফরম্যান্স সংখ্যাসূচক কম্পিউটিং লাইব্রেরি) ব্যবহার করে একটি ইংরেজি-ফরাসি অনুবাদ কাজের জন্য RecurrentGemma 2B নির্দেশ মডেলটিকে কীভাবে সূক্ষ্ম-টিউন করতে হয় তা প্রদর্শন করে। নির্ভরযোগ্য JAX কোড লেখার জন্য ইউটিলিটিগুলির একটি লাইব্রেরি), Optax (JAX-ভিত্তিক গ্রেডিয়েন্ট প্রসেসিং এবং অপ্টিমাইজেশান লাইব্রেরি), এবং MTNT (মেশিন ট্রান্সলেশন অফ নয়জি টেক্সট) ডেটাসেট । যদিও এই নোটবুকে ফ্ল্যাক্স সরাসরি ব্যবহার করা হয়নি, তবে জেমা তৈরি করতে ফ্ল্যাক্স ব্যবহার করা হয়েছিল।

recurrentgemma লাইব্রেরিটি JAX, Flax, Orbax (চেকপয়েন্টিংয়ের মতো প্রশিক্ষণের জন্য একটি JAX-ভিত্তিক লাইব্রেরি), এবং SentencePiece (একটি টোকেনাইজার/ডিটোকেনাইজার লাইব্রেরি) দিয়ে লেখা হয়েছিল।

এই নোটবুকটি T4 GPU সহ Google Colab-এ চলতে পারে ( এডিট > নোটবুক সেটিংস > হার্ডওয়্যার এক্সিলারেটরের অধীনে T4 GPU নির্বাচন করুন)।

সেটআপ

নিম্নলিখিত বিভাগগুলি একটি RecurrentGemma মডেল ব্যবহার করার জন্য একটি নোটবুক প্রস্তুত করার পদক্ষেপগুলি ব্যাখ্যা করে, যার মধ্যে মডেল অ্যাক্সেস, একটি API কী পাওয়া এবং নোটবুক রানটাইম কনফিগার করা।

জেমার জন্য কাগল অ্যাক্সেস সেট আপ করুন

এই টিউটোরিয়ালটি সম্পূর্ণ করতে, আপনাকে প্রথমে কয়েকটি ব্যতিক্রম সহ জেমা সেটআপের অনুরূপ সেটআপ নির্দেশাবলী অনুসরণ করতে হবে:

kaggle.com- এ RecurrentGemma (Gemma এর পরিবর্তে) অ্যাক্সেস পান।
RecurrentGemma মডেল চালানোর জন্য পর্যাপ্ত সম্পদ সহ একটি Colab রানটাইম বেছে নিন।
একটি Kaggle ব্যবহারকারীর নাম এবং API কী তৈরি এবং কনফিগার করুন।

আপনি RecurrentGemma সেটআপ সম্পূর্ণ করার পরে, পরবর্তী বিভাগে যান, যেখানে আপনি আপনার Colab পরিবেশের জন্য পরিবেশের ভেরিয়েবল সেট করবেন।

পরিবেশের ভেরিয়েবল সেট করুন

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')

`recurrentgemma` লাইব্রেরি ইনস্টল করুন

ফ্রি Colab হার্ডওয়্যার ত্বরণ বর্তমানে এই নোটবুক চালানোর জন্য অপর্যাপ্ত । আপনি যদি Colab Pay As You Go বা Colab Pro ব্যবহার করেন, তাহলে Edit > Notebook সেটিংস > A100 GPU বেছে নিন > হার্ডওয়্যার অ্যাক্সিলারেশন চালু করতে সেভ করুন -এ ক্লিক করুন।

এরপরে, আপনাকে github.com/google-deepmind/recurrentgemma থেকে Google DeepMind recurrentgemma লাইব্রেরি ইনস্টল করতে হবে। আপনি যদি "পিপের নির্ভরতা সমাধানকারী" সম্পর্কে একটি ত্রুটি পান তবে আপনি সাধারণত এটি উপেক্ষা করতে পারেন।

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

লাইব্রেরি আমদানি করুন

এই নোটবুকটি ফ্ল্যাক্স (নিউরাল নেটওয়ার্কের জন্য), কোর JAX , SentencePiece (টোকেনাইজেশনের জন্য), Chex (নির্ভরযোগ্য JAX কোড লেখার জন্য ইউটিলিটিগুলির একটি লাইব্রেরি), Optax (গ্রেডিয়েন্ট প্রসেসিং এবং অপ্টিমাইজেশান লাইব্রেরি), এবং TensorFlow ডেটাসেট ব্যবহার করে।

import pathlib
from typing import Any, Mapping, Iterator
import enum
import functools

import chex
import jax
import jax.numpy as jnp
import optax

import tensorflow as tf
import tensorflow_datasets as tfds

import sentencepiece as spm

from recurrentgemma import jax as recurrentgemma

RecurrentGemma মডেল লোড করুন

kagglehub.model_download দিয়ে RecurrentGemma মডেল লোড করুন, যা তিনটি আর্গুমেন্ট নেয়:

handle : কাগল থেকে মডেল হ্যান্ডেল
path : (ঐচ্ছিক স্ট্রিং) স্থানীয় পথ
force_download : (ঐচ্ছিক বুলিয়ান) মডেলটিকে পুনরায় ডাউনলোড করতে বাধ্য করে

RECURRENTGEMMA_VARIANT = '2b-it' # @param ['2b', '2b-it'] {type:"string"}

import kagglehub

RECURRENTGEMMA_PATH = kagglehub.model_download(f'google/recurrentgemma/flax/{RECURRENTGEMMA_VARIANT}')

Downloading from https://www.kaggle.com/api/v1/models/google/recurrentgemma/flax/2b-it/1/download...
100%|██████████| 3.85G/3.85G [00:50<00:00, 81.5MB/s]
Extracting model files...

print('RECURRENTGEMMA_VARIANT:', RECURRENTGEMMA_VARIANT)

RECURRENTGEMMA_VARIANT: 2b-it

মডেল ওজন এবং টোকেনাইজারের অবস্থান পরীক্ষা করুন, তারপর পাথ ভেরিয়েবল সেট করুন। টোকেনাইজার ডিরেক্টরিটি মূল ডিরেক্টরিতে থাকবে যেখানে আপনি মডেলটি ডাউনলোড করেছেন, যখন মডেলের ওজনগুলি একটি সাব-ডিরেক্টরিতে থাকবে। উদাহরণ স্বরূপ:

tokenizer.model ফাইলটি /LOCAL/PATH/TO/recurrentgemma/flax/2b-it/1 ) এ থাকবে।
মডেল চেকপয়েন্টটি হবে /LOCAL/PATH/TO/recurrentgemma/flax/2b-it/1/2b-it )।

CKPT_PATH = os.path.join(RECURRENTGEMMA_PATH, RECURRENTGEMMA_VARIANT)
TOKENIZER_PATH = os.path.join(RECURRENTGEMMA_PATH, 'tokenizer.model')
print('CKPT_PATH:', CKPT_PATH)
print('TOKENIZER_PATH:', TOKENIZER_PATH)

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

MTNT ডেটাসেট এবং জেমা টোকেনাইজার লোড করুন এবং প্রস্তুত করুন

আপনি MTNT (মেশিন ট্রান্সলেশন অফ নয়েসি টেক্সট) ডেটাসেট ব্যবহার করবেন, যা টেনসরফ্লো ডেটাসেট থেকে পাওয়া যায়।

MTNT ডেটাসেটের ইংরেজি-থেকে-ফ্রেঞ্চ ডেটাসেট অংশ ডাউনলোড করুন এবং তারপরে দুটি উদাহরণের নমুনা নিন। ডেটাসেটের প্রতিটি নমুনায় দুটি এন্ট্রি রয়েছে: src : মূল ইংরেজি বাক্য; এবং dst : সংশ্লিষ্ট ফরাসি অনুবাদ।

ds = tfds.load("mtnt/en-fr", split="train")

ds = ds.take(2)
ds = ds.as_numpy_iterator()

for idx, example in enumerate(ds):
  print(f'Example {idx}:')
  for key, val in example.items():
    print(f'{key}: {val}')
  print()

Downloading and preparing dataset 35.08 MiB (download: 35.08 MiB, generated: 11.33 MiB, total: 46.41 MiB) to /root/tensorflow_datasets/mtnt/en-fr/1.0.0...
Dl Completed...: 0 url [00:00, ? url/s]
Dl Size...: 0 MiB [00:00, ? MiB/s]
Extraction completed...: 0 file [00:00, ? file/s]
Generating splits...:   0%|          | 0/3 [00:00<?, ? splits/s]
Generating train examples...:   0%|          | 0/35692 [00:00<?, ? examples/s]
Shuffling /root/tensorflow_datasets/mtnt/en-fr/1.0.0.incompleteJLH33K/mtnt-train.tfrecord*...:   0%|          …
Generating test examples...:   0%|          | 0/1020 [00:00<?, ? examples/s]
Shuffling /root/tensorflow_datasets/mtnt/en-fr/1.0.0.incompleteJLH33K/mtnt-test.tfrecord*...:   0%|          |…
Generating valid examples...:   0%|          | 0/811 [00:00<?, ? examples/s]
Shuffling /root/tensorflow_datasets/mtnt/en-fr/1.0.0.incompleteJLH33K/mtnt-valid.tfrecord*...:   0%|          …
Dataset mtnt downloaded and prepared to /root/tensorflow_datasets/mtnt/en-fr/1.0.0. Subsequent calls will reuse this data.
Example 0:
dst: b'Le groupe de " toutes les \xc3\xa9toiles potentielles de la conf\xc3\xa9rence de l\'Est mais qui ne s\'en sortent pas dans le groupe de l\'Ouest ".'
src: b'The group of \xe2\x80\x9ceastern conference potential all stars but not making it in the West\xe2\x80\x9d group.'

Example 1:
dst: b"Kameron est-elle un peu aigrie de son manque de temps \xc3\xa0 l'\xc3\xa9cran ?"
src: b'Is Kameron a Little Salty About Her Lack of Air Time?'

sentencepiece.SentencePieceProcessor ব্যবহার করে তৈরি জেমা টোকেনাইজার লোড করুন:

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

True

ইংরেজি থেকে ফরাসি অনুবাদ টাস্কের জন্য SentencePieceProcessor কাস্টমাইজ করুন। যেহেতু আপনি RecurrentGemma (Griffin) মডেলের ইংরেজি অংশটি সূক্ষ্ম-টিউনিং করবেন, তাই আপনাকে কয়েকটি সমন্বয় করতে হবে, যেমন:

ইনপুট উপসর্গ : প্রতিটি ইনপুটে একটি সাধারণ উপসর্গ যোগ করা অনুবাদের কাজকে সংকেত দেয়। উদাহরণস্বরূপ, আপনি একটি উপসর্গ সহ একটি প্রম্পট ব্যবহার করতে পারেন যেমন Translate this into French: [INPUT_SENTENCE] ।
অনুবাদ শুরুর প্রত্যয় : প্রতিটি প্রম্পটের শেষে একটি প্রত্যয় যোগ করলে অনুবাদ প্রক্রিয়া কখন শুরু করতে হবে তা Gemma মডেলকে নির্দেশ করে। একটি নতুন লাইন কাজ করা উচিত.
ভাষার মডেল টোকেন : RecurrentGemma (Griffin) মডেল প্রতিটি সিকোয়েন্সের শুরুতে একটি "বিগিনিং অফ সিকোয়েন্স" টোকেন আশা করে। একইভাবে, আপনাকে প্রতিটি প্রশিক্ষণ উদাহরণের শেষে একটি "অনুক্রমের শেষ" টোকেন যোগ করতে হবে।

নিম্নরূপ SentencePieceProcessor এর চারপাশে একটি কাস্টম মোড়ক তৈরি করুন:

class GriffinTokenizer:
  """A custom wrapper around a SentencePieceProcessor."""

  def __init__(self, spm_processor: spm.SentencePieceProcessor):
    self._spm_processor = spm_processor

  @property
  def pad_id(self) -> int:
    """Fast access to the pad ID."""
    return self._spm_processor.pad_id()

  def tokenize(
      self,
      example: str | bytes,
      prefix: str = '',
      suffix: str = '',
      add_eos: bool = True,
  ) -> jax.Array:
    """
    A tokenization function.

    Args:
      example: Input string to tokenize.
      prefix:  Prefix to add to the input string.
      suffix:  Suffix to add to the input string.
      add_eos: If True, add an end of sentence token at the end of the output
               sequence.
    Returns:
      Tokens corresponding to the input string.
    """
    int_list = [self._spm_processor.bos_id()]
    int_list.extend(self._spm_processor.EncodeAsIds(prefix + example + suffix))
    if add_eos:
      int_list.append(self._spm_processor.eos_id())

    return jnp.array(int_list, dtype=jnp.int32)

  def tokenize_tf_op(
      self,
      str_tensor: tf.Tensor,
      prefix: str = '',
      suffix: str = '',
      add_eos: bool = True,
  ) -> tf.Tensor:
    """A TensforFlow operator for the `tokenize` function."""
    encoded = tf.numpy_function(
        self.tokenize,
        [str_tensor, prefix, suffix, add_eos],
        tf.int32)
    encoded.set_shape([None])
    return encoded

  def to_string(self, tokens: jax.Array) -> str:
    """Convert an array of tokens to a string."""
    return self._spm_processor.EncodeIds(tokens.tolist())

আপনার নতুন কাস্টম GriffinTokenizer ইনস্ট্যান্টিয়েট করে এটি ব্যবহার করে দেখুন, এবং তারপর MTNT ডেটাসেটের একটি ছোট নমুনায় এটি প্রয়োগ করুন:

def tokenize_source(tokenizer, example: tf.Tensor):
  return tokenizer.tokenize_tf_op(
      example,
      prefix='Translate this into French:\n',
      suffix='\n',
      add_eos=False
  )
def tokenize_destination(tokenizer, example: tf.Tensor):
  return tokenizer.tokenize_tf_op(example, add_eos=True)

tokenizer = GriffinTokenizer(vocab)

ds = tfds.load("mtnt/en-fr",split="train")
ds = ds.take(2)
ds = ds.map(lambda x: {
    'src': tokenize_source(tokenizer, x['src']),
    'dst': tokenize_destination(tokenizer, x['dst'])
  })
ds = ds.as_numpy_iterator()

for idx, example in enumerate(ds):
  print(f'Example {idx}:')
  for key, val in example.items():
    print(f'{key}: {val}')
  print()

Example 0:
src: [     2  49688    736   1280   6987 235292    108    651   2778    576
   1080 104745  11982   5736    832   8995    901    780   3547    665
    575    573   4589 235369   2778 235265    108]
dst: [     2   2025  29653    581    664  16298   1437  55563  41435   7840
    581    683 111452    581    533 235303   9776   4108   2459    679
    485 235303    479   6728    579   1806   2499    709  29653    581
    533 235303 101323  16054      1]

Example 1:
src: [     2  49688    736   1280   6987 235292    108   2437  87150    477
    476  11709 230461   8045   3636  40268    576   4252   4897 235336
    108]
dst: [     2 213606    477   1455 235290   3510    748   8268 191017   2809
    581   2032  69972    581  11495   1305    533 235303  65978   1654
      1]

সমগ্র MTNT ডেটাসেটের জন্য একটি ডেটা লোডার তৈরি করুন:

@chex.dataclass(frozen=True)
class TrainingInput:
  # Input tokens provided to the model.
  input_tokens: jax.Array

  # A mask that determines which tokens contribute to the target loss
  # calculation.
  target_mask: jax.Array

class DatasetSplit(enum.Enum):
  TRAIN = 'train'
  VALIDATION = 'valid'


class MTNTDatasetBuilder:
  """A data loader for the MTNT dataset."""

  N_ITEMS = {DatasetSplit.TRAIN: 35_692, DatasetSplit.VALIDATION: 811}

  BUFFER_SIZE_SHUFFLE = 10_000
  TRANSLATION_PREFIX = 'Translate this into French:\n'
  TRANSLATION_SUFFIX = '\n'

  def __init__(self,
               tokenizer : GriffinTokenizer,
               max_seq_len: int):
    """A constructor.

    Args:
      tokenizer: The tokenizer to use.
      max_seq_len: The size of each sequence in a given batch.
    """
    self._tokenizer = tokenizer
    self._base_data = {
        DatasetSplit.TRAIN: tfds.load("mtnt/en-fr",split="train"),
        DatasetSplit.VALIDATION: tfds.load("mtnt/en-fr",split="valid"),
    }
    self._max_seq_len = max_seq_len

  def _tokenize_source(self, example: tf.Tensor):
    """A tokenization function for the source."""
    return self._tokenizer.tokenize_tf_op(
        example, prefix=self.TRANSLATION_PREFIX, suffix=self.TRANSLATION_SUFFIX,
        add_eos=False
    )

  def _tokenize_destination(self, example: tf.Tensor):
    """A tokenization function for the French translation."""
    return self._tokenizer.tokenize_tf_op(example, add_eos=True)

  def _pad_up_to_max_len(self,
                         input_tensor: tf.Tensor,
                         pad_value: int | bool,
                         ) -> tf.Tensor:
    """Pad the given tensor up to sequence length of a batch."""
    seq_len = tf.shape(input_tensor)[0]
    to_pad = tf.maximum(self._max_seq_len - seq_len, 0)
    return tf.pad(
        input_tensor, [[0, to_pad]], mode='CONSTANT', constant_values=pad_value,
    )

  def _to_training_input(
      self,
      src_tokens: jax.Array,
      dst_tokens: jax.Array,
  ) -> TrainingInput:
    """Build a training input from a tuple of source and destination tokens."""

    # The input sequence fed to the model is simply the concatenation of the
    # source and the destination.
    tokens = tf.concat([src_tokens, dst_tokens], axis=0)

    # You want to prevent the model from updating based on the source (input)
    # tokens. To achieve this, add a target mask to each input.
    q_mask = tf.zeros_like(src_tokens, dtype=tf.bool)
    a_mask = tf.ones_like(dst_tokens, dtype=tf.bool)
    mask = tf.concat([q_mask, a_mask], axis=0)

    # If the output tokens sequence is smaller than the target sequence size,
    # then pad it with pad tokens.
    tokens = self._pad_up_to_max_len(tokens, self._tokenizer.pad_id)

    # You don't want to perform the backward on the pad tokens.
    mask = self._pad_up_to_max_len(mask, False)

    return TrainingInput(input_tokens=tokens, target_mask=mask)


  def get_train_dataset(self, batch_size: int, num_epochs: int):
    """Build the training dataset."""

    # Tokenize each sample.
    ds = self._base_data[DatasetSplit.TRAIN].map(
        lambda x : (self._tokenize_source(x['src']),
                    self._tokenize_destination(x['dst']))
    )

    # Convert them to training inputs.
    ds = ds.map(lambda x, y: self._to_training_input(x, y))

    # Remove the samples which are too long.
    ds = ds.filter(lambda x: tf.shape(x.input_tokens)[0] <= self._max_seq_len)

    # Shuffle the dataset.
    ds = ds.shuffle(buffer_size=self.BUFFER_SIZE_SHUFFLE)

    # Repeat if necessary.
    ds = ds.repeat(num_epochs)

    # Build batches.
    ds = ds.batch(batch_size, drop_remainder=True)
    return ds

  def get_validation_dataset(self, batch_size: int):
    """Build the validation dataset."""

    # Same as the training dataset, but no shuffling and no repetition
    ds = self._base_data[DatasetSplit.VALIDATION].map(
        lambda x : (self._tokenize_source(x['src']),
                    self._tokenize_destination(x['dst']))
    )
    ds = ds.map(lambda x, y: self._to_training_input(x, y))
    ds = ds.filter(lambda x: tf.shape(x.input_tokens)[0] <= self._max_seq_len)
    ds = ds.batch(batch_size, drop_remainder=True)
    return ds

MTNTDatasetBuilder ব্যবহার করে দেখুন কাস্টম GriffinTokenizer আবার চালু করে, তারপর এটি MTNT ডেটাসেটে প্রয়োগ করুন এবং দুটি উদাহরণের নমুনা নিন:

dataset_builder = MTNTDatasetBuilder(tokenizer, max_seq_len=20)
ds = dataset_builder.get_train_dataset(3, 1)
ds = ds.take(2)
ds = ds.as_numpy_iterator()

for idx, example in enumerate(ds):
  print(f'Example {idx}:')
  for key, val in example.items():
    print(f'{key}: {val}')
  print()

WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class '__main__.TrainingInput'>
WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class '__main__.TrainingInput'>
WARNING:tensorflow:Mapping types may not work well with tf.nest. Prefer using MutableMapping for <class '__main__.TrainingInput'>
Example 0:
input_tokens: [[     2  49688    736   1280   6987 235292    108  12583    665 235265
     108      2   6151  94975   1320   6238 235265      1      0      0]
 [     2  49688    736   1280   6987 235292    108   4899  29960  11270
  108282 235265    108      2   4899  79025  11270 108282      1      0]
 [     2  49688    736   1280   6987 235292    108  26620 235265    108
       2  26620 235265      1      0      0      0      0      0      0]]
target_mask: [[False False False False False False False False False False False  True
   True  True  True  True  True  True False False]
 [False False False False False False False False False False False False
  False  True  True  True  True  True  True False]
 [False False False False False False False False False False  True  True
   True  True False False False False False False]]

Example 1:
input_tokens: [[     2  49688    736   1280   6987 235292    108    527   5174   1683
  235336    108      2 206790    581  20726    482   2208   1654      1]
 [     2  49688    736   1280   6987 235292    108  28484 235256 235336
     108      2 120500  13832   1654      1      0      0      0      0]
 [     2  49688    736   1280   6987 235292    108 235324 235304   2705
  235265    108      2 235324 235304  19963 235265      1      0      0]]
target_mask: [[False False False False False False False False False False False False
   True  True  True  True  True  True  True  True]
 [False False False False False False False False False False False  True
   True  True  True  True False False False False]
 [False False False False False False False False False False False False
   True  True  True  True  True  True False False]]

মডেল কনফিগার করুন

আপনি Gemma মডেলের ফাইন-টিউনিং শুরু করার আগে, আপনাকে এটি কনফিগার করতে হবে।

recurrentgemma.jax.utils.load_parameters পদ্ধতির সাথে RecurrentGemma (Griffin) মডেল চেকপয়েন্ট লোড করুন:

params =  recurrentgemma.load_parameters(CKPT_PATH, "single_device")

RecurrentGemma মডেল চেকপয়েন্ট থেকে স্বয়ংক্রিয়ভাবে সঠিক কনফিগারেশন লোড করতে, recurrentgemma.GriffinConfig.from_flax_params_or_variables ব্যবহার করুন :

config = recurrentgemma.GriffinConfig.from_flax_params_or_variables(params)

recurrentgemma.jax.Griffin এর মাধ্যমে গ্রিফিন মডেলটি চালু করুন:

model = recurrentgemma.Griffin(config)

recurrentgemma.jax.Sampler এর উপরে RecurrentGemma মডেলের চেকপয়েন্ট/ওজন এবং টোকেনাইজার দিয়ে একটি sampler তৈরি করুন যাতে আপনার মডেল অনুবাদ করতে পারে কিনা:

sampler = recurrentgemma.Sampler(model=model, vocab=vocab, params=params)

মডেলটি ফাইন-টিউন করুন

এই বিভাগে, আপনি করবেন:

ফরোয়ার্ড পাস এবং লস ফাংশন তৈরি করতে gemma.transformer.Transformer ক্লাস ব্যবহার করুন।
টোকেনগুলির জন্য অবস্থান এবং মনোযোগ মাস্ক ভেক্টর তৈরি করুন
Flax এর সাথে একটি প্রশিক্ষণ ধাপ ফাংশন তৈরি করুন।
পিছনের পাস ছাড়াই বৈধতা ধাপ তৈরি করুন।
প্রশিক্ষণ লুপ তৈরি করুন।
জেমা মডেলটি সূক্ষ্ম সুর করুন।

recurrentgemma.jax.griffin.Griffin ক্লাস ব্যবহার করে ফরওয়ার্ড পাস এবং লস ফাংশন সংজ্ঞায়িত করুন। RecurrentGemma Griffin flax.linen.Module থেকে উত্তরাধিকারসূত্রে প্রাপ্ত, এবং দুটি প্রয়োজনীয় পদ্ধতি অফার করে:

init : মডেলের পরামিতি শুরু করে।
apply : মডেলের __call__ ফাংশনটি প্রদত্ত প্যারামিটারের সেট ব্যবহার করে চালায়।

যেহেতু আপনি প্রাক-প্রশিক্ষিত জেমা ওজন নিয়ে কাজ করছেন, তাই আপনাকে init ফাংশন ব্যবহার করতে হবে না।

def forward_and_loss_fn(
    params,
    *,
    model: recurrentgemma.Griffin,
    input_tokens: jax.Array,            # Shape [B, L]
    input_mask: jax.Array,              # Shape [B, L]
    positions: jax.Array,               # Shape [B, L]
) -> jax.Array:
  """Forward pass and loss function.

  Args:
    params: model's input parameters.
    model: Griffin model to call.
    input_tokens: input tokens sequence, shape [B, L].
    input_mask: tokens to ignore when computing the loss, shape [B, L].
    positions: relative position of each token, shape [B, L].

  Returns:
    Softmax cross-entropy loss for the next-token prediction task.
  """
  batch_size = input_tokens.shape[0]
  # Forward pass on the input data.
  # No attention cache is needed here.
  # Exclude the last step as it does not appear in the targets.
  logits, _ = model.apply(
        {"params": params},
        tokens=input_tokens[:, :-1],
        segment_pos=positions[:, :-1],
        cache=None,
    )

  # Similarly, the first token cannot be predicteds.
  target_tokens = input_tokens[:, 1:]
  target_mask = input_mask[:, 1:]

  # Convert the target labels into one-hot encoded vectors.
  one_hot = jax.nn.one_hot(target_tokens, logits.shape[-1])

  # Don't update on unwanted tokens.
  one_hot = one_hot * target_mask.astype(one_hot.dtype)[...,None]

  # Normalization factor.
  norm_factor = batch_size * (jnp.sum(target_mask) + 1e-8)

  # Return the negative log-likelihood loss (NLL) function.
  return -jnp.sum(jax.nn.log_softmax(logits) * one_hot) / norm_factor

train_step ফাংশন তৈরি করুন যা ব্যাকওয়ার্ড পাস সম্পাদন করে এবং সেই অনুযায়ী মডেলের প্যারামিটার আপডেট করে, যেখানে:

jax.value_and_grad হল ফরওয়ার্ড এবং ব্যাকওয়ার্ড পাসের সময় লস ফাংশন এবং গ্রেডিয়েন্ট মূল্যায়ন করার জন্য।
optax.apply_updates প্যারামিটার আপডেট করার জন্য।

Params = Mapping[str, Any]

def get_positions(example: jax.Array, pad_id : int) -> jax.Array:
  """Builds the position vector from the given tokens."""
  pad_mask = example != pad_id
  positions = jnp.cumsum(pad_mask, axis=-1)
  # Subtract one for all positions from the first valid one as they are
  # 0-indexed
  positions = positions - (positions >= 1)
  return positions

@functools.partial(
    jax.jit,
    static_argnames=['model', 'optimizer'],
    donate_argnames=['params', 'opt_state'],
)
def train_step(
    model: recurrentgemma.Griffin,
    params: Params,
    optimizer: optax.GradientTransformation,
    opt_state: optax.OptState,
    pad_id: int,
    example: TrainingInput,
) -> tuple[jax.Array, Params, optax.OptState]:
  """The train step.

  Args:
    model: The RecurrentGemma (Griffin) model.
    params: The model's input parameters.
    optimizer: The Optax optimizer to use.
    opt_state: The input optimizer's state.
    pad_id: The ID of the pad token.
    example: The input batch.

  Returns:
    Training loss, updated parameters, updated optimizer state.
  """

  positions = get_positions(example.input_tokens, pad_id)

  # Forward and backward passes.
  train_loss, grads = jax.value_and_grad(forward_and_loss_fn)(
      params,
      model=model,
      input_tokens=example.input_tokens,
      input_mask=example.target_mask,
      positions=positions,
  )
  # Update the parameters.
  updates, opt_state = optimizer.update(grads, opt_state, params)
  params = optax.apply_updates(params, updates)

  return train_loss, params, opt_state

ব্যাকওয়ার্ড পাস ছাড়াই validation_step ফাংশন তৈরি করুন:

@functools.partial(jax.jit, static_argnames=['model'])
def validation_step(
    model: recurrentgemma.Griffin,
    params: Params,
    pad_id: int,
    example: TrainingInput,
) -> jax.Array:
  return forward_and_loss_fn(
      params,
      model=model,
      input_tokens=example.input_tokens,
      input_mask=example.target_mask,
      positions=get_positions(example.input_tokens, pad_id),
  )

প্রশিক্ষণ লুপ সংজ্ঞায়িত করুন:

def train_loop(
    model: recurrentgemma.Griffin,
    params: Params,
    optimizer: optax.GradientTransformation,
    train_ds: Iterator[TrainingInput],
    validation_ds: Iterator[TrainingInput],
    num_steps: int | None = None,
    eval_every_n: int = 20,
):
  opt_state = jax.jit(optimizer.init)(params)

  step_counter = 0
  avg_loss=0

  # The first round of the validation loss.
  n_steps_eval = 0
  eval_loss = 0
  for val_example in validation_ds.as_numpy_iterator():
    eval_loss += validation_step(
        model, params, dataset_builder._tokenizer.pad_id, val_example
    )
    n_steps_eval += 1
  print(f"Start, validation loss: {eval_loss/n_steps_eval}")

  for train_example in train_ds:
    train_loss, params, opt_state = train_step(
        model=model,
        params=params,
        optimizer=optimizer,
        opt_state=opt_state,
        pad_id=dataset_builder._tokenizer.pad_id,
        example=train_example,
    )

    step_counter += 1
    avg_loss += train_loss
    if step_counter % eval_every_n == 0:
      eval_loss = 0

      n_steps_eval = 0
      val_iterator = validation_ds.as_numpy_iterator()
      for val_example in val_iterator:
        eval_loss += validation_step(
            model,
            params,
            dataset_builder._tokenizer.pad_id,
            val_example,
        )
        n_steps_eval +=1
      avg_loss /= eval_every_n
      eval_loss /= n_steps_eval
      print(f"STEP {step_counter} training loss: {avg_loss} - eval loss: {eval_loss}")
      avg_loss=0
    if num_steps is not None and step_counter > num_steps:
      break
  return params

এখানে আপনাকে একটি (Optax) অপ্টিমাইজার বেছে নিতে হবে। ছোট মেমরি সহ ডিভাইসগুলির জন্য, আপনার SGD ব্যবহার করা উচিত, কারণ এতে অনেক কম মেমরি ফুটপ্রিন্ট রয়েছে। সেরা ফাইন-টিউনিং পারফরম্যান্স অর্জন করতে, অ্যাডাম-ডব্লিউ ব্যবহার করে দেখুন। এই নোটবুকের নির্দিষ্ট কাজের জন্য প্রতিটি অপ্টিমাইজারের জন্য সর্বোত্তম হাইপারপ্যারামিটার 2b-it চেকপয়েন্টের জন্য এই উদাহরণে প্রদান করা হয়েছে।

def griffin_weight_decay_mask(params_like: optax.Params) -> Any:
  # Don't put weight decay on the RGLRU, the embeddings and any biases
  def enable_weight_decay(path: list[Any], _: Any) -> bool:
    # Parameters in the LRU and embedder
    path = [dict_key.key for dict_key in path]
    if 'rg_lru' in path or 'embedder' in path:
      return False
    # All biases and scales
    if path[-1] in ('b', 'scale'):
      return False
    return True

  return jax.tree_util.tree_map_with_path(enable_weight_decay, params_like)

optimizer_choice = "sgd"

if optimizer_choice == "sgd":
  optimizer = optax.sgd(learning_rate=1e-3)
  num_steps = 300
elif optimizer_choice == "adamw":
  optimizer = optax.adamw(
        learning_rate=1e-4,
        b2=0.96,
        eps=1e-8,
        weight_decay=0.1,
        mask=griffin_weight_decay_mask,
    )
  num_steps = 100
else:
  raise ValueError(f"Unknown optimizer: {optimizer_choice}")

প্রশিক্ষণ এবং বৈধতা ডেটাসেট প্রস্তুত করুন:

# Choose a small sequence length size, so that everything fits in memory.
num_epochs = 1
batch_size = 1
sequence_length = 32

# Make the dataset builder.
tokenizer = GriffinTokenizer(vocab)
dataset_builder= MTNTDatasetBuilder(tokenizer, sequence_length + 1)

# Build the training dataset.
train_ds = dataset_builder.get_train_dataset(
    batch_size=batch_size,
    num_epochs=num_epochs,
).as_numpy_iterator()

# Build the validation dataset, with a limited number of samples for this demo.
validation_ds = dataset_builder.get_validation_dataset(
    batch_size=batch_size,
).take(50)

সীমিত সংখ্যক ধাপে RecurrentGemma (Griffin) মডেলের ফাইন-টিউনিং শুরু করুন ( num_steps ):

trained_params = train_loop(
    model=model,
    params=params,
    optimizer=optimizer,
    train_ds=train_ds,
    validation_ds=validation_ds,
    num_steps=num_steps,
)

Start, validation loss: 7.894117832183838
/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/mlir.py:920: UserWarning: Some donated buffers were not usable: ShapedArray(int32[1,33]), ShapedArray(bool[1,33]), ShapedArray(int32[], weak_type=True).
See an explanation at https://jax.readthedocs.io/en/latest/faq.html#buffer_donation.
  warnings.warn("Some donated buffers were not usable:"
STEP 20 training loss: 4.592616081237793 - eval loss: 2.847407102584839
STEP 40 training loss: 2.7537424564361572 - eval loss: 2.9258534908294678
STEP 60 training loss: 2.835618257522583 - eval loss: 2.4382340908050537
STEP 80 training loss: 2.6322107315063477 - eval loss: 2.3696839809417725
STEP 100 training loss: 1.8703256845474243 - eval loss: 2.355681896209717
STEP 120 training loss: 2.7280433177948 - eval loss: 2.4059958457946777
STEP 140 training loss: 2.3047447204589844 - eval loss: 2.083082914352417
STEP 160 training loss: 2.3432137966156006 - eval loss: 2.095074415206909
STEP 180 training loss: 2.1081202030181885 - eval loss: 2.006460189819336
STEP 200 training loss: 2.5359647274017334 - eval loss: 1.9667452573776245
STEP 220 training loss: 2.202195644378662 - eval loss: 1.9440618753433228
STEP 240 training loss: 2.756615400314331 - eval loss: 2.1073737144470215
STEP 260 training loss: 2.5128934383392334 - eval loss: 2.117241859436035
STEP 280 training loss: 2.73045015335083 - eval loss: 1.9159646034240723
STEP 300 training loss: 2.0918595790863037 - eval loss: 1.9742532968521118

প্রতিটি ধাপের গণনার সাথে প্রশিক্ষণের ক্ষতি এবং বৈধতা ক্ষতি উভয়ই কমে যাওয়া উচিত।

আপনার ইনপুট প্রশিক্ষণ বিন্যাসের সাথে মেলে তা নিশ্চিত করতে, উপসর্গটি ব্যবহার করতে মনে রাখবেন Translate this into French:\n এবং শেষে একটি নতুন লাইন অক্ষর৷ এটি অনুবাদ শুরু করার জন্য মডেলটিকে সংকেত দেয়।

sampler.params = trained_params
output = sampler(
    ["Translate this into French:\nHello, my name is Morgane.\n"],
    total_generation_steps=100,
)
print(output.text[0])

/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/mlir.py:920: UserWarning: Some donated buffers were not usable: ShapedArray(int32[1,16]).
See an explanation at https://jax.readthedocs.io/en/latest/faq.html#buffer_donation.
  warnings.warn("Some donated buffers were not usable:"
Mais je m'appelle Morgane.

আরও জানুন

আপনি GitHub-এ Google DeepMind recurrentgemma লাইব্রেরি সম্পর্কে আরও জানতে পারেন, যেটিতে আপনি এই টিউটোরিয়ালে ব্যবহার করেছেন পদ্ধতি এবং মডিউলগুলির ডকস্ট্রিংগুলি রয়েছে যেমন recurrentgemma.jax.load_parameters , recurrentgemma.jax.Griffin , এবং recurrentgemma.jax.Sampler .
নিম্নলিখিত লাইব্রেরিগুলির নিজস্ব ডকুমেন্টেশন সাইট রয়েছে: কোর JAX , Flax , Chex , Optax , এবং Orbax ৷
sentencepiece টোকেনাইজার/ডিটোকেনাইজার ডকুমেন্টেশনের জন্য, Google-এর sentencepiece GitHub রেপো দেখুন।
kagglehub ডকুমেন্টেশনের জন্য, Kaggle এর kagglehub GitHub রেপোতে README.md দেখুন।
Google ক্লাউড ভার্টেক্স এআই-এর সাথে জেমা মডেলগুলি কীভাবে ব্যবহার করবেন তা শিখুন।
আপনি যদি Google ক্লাউড TPUs (v3-8 এবং নতুন) ব্যবহার করেন, তাহলে সর্বশেষ jax[tpu] প্যাকেজেও আপডেট করতে ভুলবেন না ( !pip install -U jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html ), রানটাইম পুনরায় আরম্ভ করুন, এবং jax এবং jaxlib সংস্করণ মিলছে কিনা তা পরীক্ষা করুন ( !pip list | grep jax )। এটি RuntimeError প্রতিরোধ করতে পারে যা jaxlib এবং jax সংস্করণের অমিলের কারণে হতে পারে। আরও JAX ইনস্টলেশন নির্দেশাবলীর জন্য, JAX ডক্স দেখুন।
RecurrentGemma দেখুন: Google DeepMind-এর দক্ষ ওপেন ল্যাঙ্গুয়েজ মডেলের কাগজের জন্য মুভিং পাস্ট ট্রান্সফরমার।
Griffin পড়ুন: RecurrentGemma দ্বারা ব্যবহৃত মডেল আর্কিটেকচার সম্পর্কে আরও জানতে Google DeepMind-এর দক্ষ ভাষা মডেলের কাগজের জন্য স্থানীয় মনোযোগের সাথে গেটেড লিনিয়ার পুনরাবৃত্তির মিশ্রণ।