formatting

Author: davidADSP

notebooks/09_transformer/gpt/gpt.ipynb

@@ -120,7 +120,14 @@
    "source": [
     "# Filter the dataset\n",
     "filtered_data = [\n",
-    "    \"wine review : \" + x[\"country\"] + \" : \" + x[\"province\"] + \" : \" + x[\"variety\"] + \" : \" + x[\"description\"]\n",
+    "    \"wine review : \"\n",
+    "    + x[\"country\"]\n",
+    "    + \" : \"\n",
+    "    + x[\"province\"]\n",
+    "    + \" : \"\n",
+    "    + x[\"variety\"]\n",
+    "    + \" : \"\n",
+    "    + x[\"description\"]\n",
     "    for x in wine_data\n",
     "    if x[\"country\"] is not None\n",
     "    and x[\"province\"] is not None\n",
@@ -203,7 +210,11 @@
    "outputs": [],
    "source": [
     "# Convert to a Tensorflow Dataset\n",
-    "text_ds = tf.data.Dataset.from_tensor_slices(text_data).batch(BATCH_SIZE).shuffle(1000)"
+    "text_ds = (\n",
+    "    tf.data.Dataset.from_tensor_slices(text_data)\n",
+    "    .batch(BATCH_SIZE)\n",
+    "    .shuffle(1000)\n",
+    ")"
    ]
   },
   {
@@ -340,7 +351,9 @@
     "    m = i >= j - n_src + n_dest\n",
     "    mask = tf.cast(m, dtype)\n",
     "    mask = tf.reshape(mask, [1, n_dest, n_src])\n",
-    "    mult = tf.concat([tf.expand_dims(batch_size, -1), tf.constant([1, 1], dtype=tf.int32)], 0)\n",
+    "    mult = tf.concat(\n",
+    "        [tf.expand_dims(batch_size, -1), tf.constant([1, 1], dtype=tf.int32)], 0\n",
+    "    )\n",
     "    return tf.tile(mask, mult)\n",
     "\n",
     "\n",
@@ -370,7 +383,9 @@
     "        self.embed_dim = embed_dim\n",
     "        self.ff_dim = ff_dim\n",
     "        self.dropout_rate = dropout_rate\n",
-    "        self.attn = layers.MultiHeadAttention(num_heads, key_dim, output_shape=embed_dim)\n",
+    "        self.attn = layers.MultiHeadAttention(\n",
+    "            num_heads, key_dim, output_shape=embed_dim\n",
+    "        )\n",
     "        self.dropout_1 = layers.Dropout(self.dropout_rate)\n",
     "        self.ln_1 = layers.LayerNormalization(epsilon=1e-6)\n",
     "        self.ffn_1 = layers.Dense(self.ff_dim, activation=\"relu\")\n",
@@ -382,9 +397,14 @@
     "        input_shape = tf.shape(inputs)\n",
     "        batch_size = input_shape[0]\n",
     "        seq_len = input_shape[1]\n",
-    "        causal_mask = causal_attention_mask(batch_size, seq_len, seq_len, tf.bool)\n",
+    "        causal_mask = causal_attention_mask(\n",
+    "            batch_size, seq_len, seq_len, tf.bool\n",
+    "        )\n",
     "        attention_output, attention_scores = self.attn(\n",
-    "            inputs, inputs, attention_mask=causal_mask, return_attention_scores=True\n",
+    "            inputs,\n",
+    "            inputs,\n",
+    "            attention_mask=causal_mask,\n",
+    "            return_attention_scores=True,\n",
     "        )\n",
     "        attention_output = self.dropout_1(attention_output)\n",
     "        out1 = self.ln_1(inputs + attention_output)\n",
@@ -430,7 +450,9 @@
     "        self.max_len = max_len\n",
     "        self.vocab_size = vocab_size\n",
     "        self.embed_dim = embed_dim\n",
-    "        self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)\n",
+    "        self.token_emb = layers.Embedding(\n",
+    "            input_dim=vocab_size, output_dim=embed_dim\n",
+    "        )\n",
     "        self.pos_emb = layers.Embedding(input_dim=max_len, output_dim=embed_dim)\n",
     "\n",
     "    def call(self, x):\n",
@@ -469,7 +491,9 @@
    "source": [
     "inputs = layers.Input(shape=(None,), dtype=tf.int32)\n",
     "x = TokenAndPositionEmbedding(MAX_LEN, VOCAB_SIZE, EMBEDDING_DIM)(inputs)\n",
-    "x, attention_scores = TransformerBlock(N_HEADS, KEY_DIM, EMBEDDING_DIM, FEED_FORWARD_DIM)(x)\n",
+    "x, attention_scores = TransformerBlock(\n",
+    "    N_HEADS, KEY_DIM, EMBEDDING_DIM, FEED_FORWARD_DIM\n",
+    ")(x)\n",
     "outputs = layers.Dense(VOCAB_SIZE, activation=\"softmax\")(x)\n",
     "gpt = models.Model(inputs=inputs, outputs=[outputs, attention_scores])\n",
     "gpt.compile(\"adam\", loss=[losses.SparseCategoricalCrossentropy(), None])"
@@ -518,22 +542,32 @@
     "class TextGenerator(callbacks.Callback):\n",
     "    def __init__(self, index_to_word, top_k=10):\n",
     "        self.index_to_word = index_to_word\n",
-    "        self.word_to_index = {word: index for index, word in enumerate(index_to_word)}\n",
+    "        self.word_to_index = {\n",
+    "            word: index for index, word in enumerate(index_to_word)\n",
+    "        }\n",
     "\n",
     "    def sample_from(self, probs, temperature):\n",
     "        probs = probs ** (1 / temperature)\n",
     "        probs = probs / np.sum(probs)\n",
     "        return np.random.choice(len(probs), p=probs), probs\n",
     "\n",
     "    def generate(self, start_prompt, max_tokens, temperature):\n",
-    "        start_tokens = [self.word_to_index.get(x, 1) for x in start_prompt.split()]\n",
+    "        start_tokens = [\n",
+    "            self.word_to_index.get(x, 1) for x in start_prompt.split()\n",
+    "        ]\n",
     "        sample_token = None\n",
     "        info = []\n",
     "        while len(start_tokens) < max_tokens and sample_token != 0:\n",
     "            x = np.array([start_tokens])\n",
     "            y, att = self.model.predict(x, verbose=0)\n",
     "            sample_token, probs = self.sample_from(y[0][-1], temperature)\n",
-    "            info.append({\"prompt\": start_prompt, \"word_probs\": probs, \"atts\": att[0, :, -1, :]})\n",
+    "            info.append(\n",
+    "                {\n",
+    "                    \"prompt\": start_prompt,\n",
+    "                    \"word_probs\": probs,\n",
+    "                    \"atts\": att[0, :, -1, :],\n",
+    "                }\n",
+    "            )\n",
     "            start_tokens.append(sample_token)\n",
     "            start_prompt = start_prompt + \" \" + self.index_to_word[sample_token]\n",
     "        print(f\"\\ngenerated text:\\n{start_prompt}\\n\")\n",
@@ -611,7 +645,9 @@
     "def print_probs(info, vocab, top_k=5):\n",
     "    for i in info:\n",
     "        highlighted_text = []\n",
-    "        for word, att_score in zip(i[\"prompt\"].split(), np.mean(i[\"atts\"], axis=0)):\n",
+    "        for word, att_score in zip(\n",
+    "            i[\"prompt\"].split(), np.mean(i[\"atts\"], axis=0)\n",
+    "        ):\n",
     "            highlighted_text.append(\n",
     "                '<span style=\"background-color:rgba(135,206,250,'\n",
     "                + str(att_score / max(np.mean(i[\"atts\"], axis=0)))\n",
@@ -637,7 +673,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "info = text_generator.generate(\"wine review : us\", max_tokens=80, temperature=1.0)"
+    "info = text_generator.generate(\n",
+    "    \"wine review : us\", max_tokens=80, temperature=1.0\n",
+    ")"
    ]
   },
   {
@@ -647,7 +685,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "info = text_generator.generate(\"wine review : italy\", max_tokens=80, temperature=0.5)"
+    "info = text_generator.generate(\n",
+    "    \"wine review : italy\", max_tokens=80, temperature=0.5\n",
+    ")"
    ]
   },
   {
@@ -657,7 +697,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "info = text_generator.generate(\"wine review : germany\", max_tokens=80, temperature=0.5)\n",
+    "info = text_generator.generate(\n",
+    "    \"wine review : germany\", max_tokens=80, temperature=0.5\n",
+    ")\n",
     "print_probs(info, vocab)"
    ]
   },