Upload Phi2.py

Phi2.py

@@ -0,0 +1,196 @@
+import tensorflow as tf
+from tensorflow.keras.layers import Dense,LayerNormalization,Embedding
+from tensorflow.keras import Model
+import math
+from dataclasses import dataclass
+
+
+@dataclass
+class ModelArgs:
+    n_positions: int = 2048
+    vocab_size: int = 51200
+    n_embd: int = 2560
+    n_head: int = 32
+    n_layer: int = 32
+    rotary_dim: int = 32
+
+
+class RoPEAttention:
+    def __init__(self, dims: int, n_head: int, rotary_dim: int):
+        self.n_head = n_head
+
+        self.q_proj = Dense(dims)
+        self.k_proj = Dense(dims)
+        self.v_proj = Dense(dims)
+        self.dense = Dense(dims)
+
+        self.rope = RoPE(rotary_dim, traditional=False)
+
+    def __call__(self, x, mask=None, cache=None):
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        # Extract some shapes
+        n_head = self.n_head
+        B, L, D = queries.shape
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = tf.transpose(tf.reshape(queries, (B, L, n_head, -1)), (0, 2, 1, 3))
+        keys = tf.transpose(tf.reshape(keys, (B, L, n_head, -1)), (0, 2, 1, 3))
+        values = tf.transpose(tf.reshape(values, (B, L, n_head, -1)), (0, 2, 1, 3))
+
+        # Add RoPE to the queries and keys and combine them with the cache
+        if cache is not None:
+            key_cache, value_cache = cache
+            queries = self.rope(queries, offset=key_cache.shape[2])
+            keys = self.rope(keys, offset=key_cache.shape[2])
+            keys = tf.concat([key_cache, keys], axis=2)
+            values = tf.concat([value_cache, values], axis=2)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        queries = tf.cast(queries, tf.float32)
+        keys = tf.cast(keys, tf.float32)
+
+        # Finally perform the attention computation
+        scale = math.sqrt(1 / queries.shape[-1])
+        scores = tf.matmul((queries * scale), tf.transpose(keys, (0, 1, 3, 2)))
+        if mask is not None:
+            scores = scores + mask
+
+        scores = tf.cast(tf.nn.softmax(scores, axis=-1), values.dtype)
+        values_hat = tf.reshape(tf.transpose(tf.matmul(scores, values), (0, 2, 1, 3)), (B, L, -1))
+
+        return self.dense(values_hat), (keys, values)
+
+
+class MLP:
+    def __init__(self, dim, hidden_dim):
+        self.fc1 = Dense(hidden_dim)
+        self.fc2 = Dense(dim)
+
+    def __call__(self, x):
+        return self.fc2(tf.nn.gelu(self.fc1(x), approximate="precise"))
+
+
+class ParallelBlock:
+    def __init__(self, config: ModelArgs):
+        dims = config.n_embd
+        mlp_dims = dims * 4
+        self.self_attn = RoPEAttention(dims, config.n_head, config.rotary_dim)
+        self.input_layernorm = LayerNormalization()
+        self.mlp = MLP(dims, mlp_dims)
+
+    def __call__(self, x, mask, cache):
+        h = self.input_layernorm(x)
+        attn_h, cache = self.self_attn(h, mask, cache)
+        ff_h = self.mlp(h)
+        return attn_h + ff_h + x, cache
+
+
+class Transformer:
+    def __init__(self, config: ModelArgs):
+        self.embed_tokens = Embedding(config.vocab_size, config.n_embd)
+        self.layers = [ParallelBlock(config) for i in range(config.n_layer)]
+        self.final_layernorm = LayerNormalization()
+
+    def __call__(self, x, mask, cache):
+        x = self.embed_tokens(x)
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for e, layer in enumerate(self.layers):
+            x, cache[e] = layer(x, mask, cache[e])
+        return self.final_layernorm(x), cache
+
+
+class Phi2(Model):
+    def __init__(self, config: ModelArgs):
+        super(Phi2, self).__init__()
+        self.model = Transformer(config)
+        self.lm_head = Dense(config.vocab_size)
+
+    def __call__(
+        self,
+        x,
+        mask = None,
+        cache = None,
+    ):
+        mask = None
+        if x.shape[1] > 1:
+            mask = tf.fill((x.shape[1], x.shape[1]), float("-inf"))
+            mask = tf.linalg.band_part(mask, 0, -1)
+            mask = tf.linalg.set_diag(mask, tf.zeros(x.shape[1]))
+            mask = tf.cast(mask, x.dtype)
+
+        y, cache = self.model(x, mask, cache)
+        return self.lm_head(y), cache
+
+
+class RoPE:
+    def __init__(self, dims: int, traditional: bool = False, base=None):
+        self.dims = dims
+        self.traditional = traditional
+        self.base = base
+
+    def _compute_rope(self, costheta, sintheta, x):
+        x1 = x[..., : self.dims // 2]
+        x2 = x[..., self.dims // 2 : self.dims]
+        rx1 = x1 * costheta - x2 * sintheta
+        rx2 = x1 * sintheta + x2 * costheta
+
+        if self.dims < x.shape[-1]:
+            rx = tf.concat([rx1, rx2, x[..., self.dims :]], axis=-1)
+        else:
+            rx = tf.concat([rx1, rx2], axis=-1)
+
+        return rx
+
+    def _compute_traditional_rope(self, costheta, sintheta, x):
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+        rx1 = x1 * costheta - x2 * sintheta
+        rx2 = x1 * sintheta + x2 * costheta
+
+        if self.dims < x.shape[-1]:
+            raise NotImplementedError(
+                "RoPE doesn't implement partial traditional application"
+            )
+
+        rx = tf.concat([rx1[..., None], rx2[..., None]], axis=-1)
+
+        return rx
+
+    def __call__(self, x, offset: int = 0):
+        shape = x.shape
+        x = tf.reshape(x, (-1, shape[-2], shape[-1]))
+        N = x.shape[1] + offset
+        costheta, sintheta = RoPE.create_cos_sin_theta(
+            N, self.dims, offset=offset, base=self.base, dtype=x.dtype
+        )
+
+        rope = (
+            self._compute_traditional_rope if self.traditional else self._compute_rope
+        )
+        rx = rope(costheta, sintheta, x)
+
+        return tf.reshape(rx, shape)
+
+    @staticmethod
+    def create_cos_sin_theta(
+        N: int,
+        D: int,
+        offset: int = 0,
+        base: float = 10000,
+        dtype=tf.float32,
+    ):
+        D = D // 2
+        positions = tf.range(offset, N, dtype=dtype)
+        freqs = tf.math.exp(
+            -tf.range(0, D, dtype=dtype) * (tf.math.log(base) / D)
+        )
+        theta = tf.reshape(positions, (-1, 1)) * tf.reshape(freqs, (1, -1))
+        costheta = tf.math.cos(theta)
+        sintheta = tf.math.sin(theta)
+
+        return costheta, sintheta