Upload indigotf/model.py with huggingface_hub

indigotf/model.py

@@ -43,19 +43,24 @@ def build_gpt(vocab_size, block_size, n_layer=4, n_head=4, n_embd=128, dropout=0
     return tf.keras.Model(tokens, logits, name=name)
 
 
+@tf.function(reduce_retracing=True)
+def _langkah(model, idx_cond, temperature, top_k):
+    logits = model(idx_cond, training=False)[:, -1, :]
+    logits = logits / max(temperature, 1e-8)
+    if top_k is not None:
+        k = min(top_k, int(logits.shape[-1]))
+        vals, _ = tf.math.top_k(logits, k=k)
+        logits = tf.where(
+            logits < vals[:, -1:],
+            tf.fill(tf.shape(logits), tf.float32.min),
+            logits,
+        )
+    return tf.random.categorical(logits, num_samples=1, dtype=tf.int64)
+
+
 def generate(model, idx, max_new_tokens, block_size, temperature=1.0, top_k=None):
     for _ in range(max_new_tokens):
         idx_cond = idx[:, -block_size:]
-        logits = model(idx_cond, training=False)[:, -1, :]
-        logits = logits / max(temperature, 1e-8)
-        if top_k is not None:
-            k = min(top_k, int(logits.shape[-1]))
-            vals, _ = tf.math.top_k(logits, k=k)
-            logits = tf.where(
-                logits < vals[:, -1:],
-                tf.fill(tf.shape(logits), tf.float32.min),
-                logits,
-            )
-        next_id = tf.random.categorical(logits, num_samples=1, dtype=tf.int64)
+        next_id = _langkah(model, idx_cond, temperature, top_k)
         idx = tf.concat([idx, next_id], axis=1)
     return idx