@zuvia-software-solutions/code-mapper 2.4.0 → 2.5.0

package/models/jina-code-0.5b-mlx/model.py

@@ -1,127 +0,0 @@
-
-from dataclasses import dataclass
-from typing import Optional, List
-import mlx.core as mx
-import mlx.nn as nn
-
-@dataclass
-class ModelArgs:
-    hidden_size: int
-    num_hidden_layers: int
-    intermediate_size: int
-    num_attention_heads: int
-    rms_norm_eps: float
-    vocab_size: int
-    num_key_value_heads: int
-    max_position_embeddings: int
-    rope_theta: float = 1000000.0
-    tie_word_embeddings: bool = True
-
-class Attention(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        dim = args.hidden_size
-        self.n_heads = args.num_attention_heads
-        self.n_kv_heads = args.num_key_value_heads
-        self.head_dim = dim // self.n_heads
-        self.scale = self.head_dim ** -0.5
-        self.rope_theta = args.rope_theta
-        self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=True)
-        self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=True)
-        self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=True)
-        self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False)
-
-    def __call__(self, x, mask=None):
-        B, L, D = x.shape
-        q = self.q_proj(x).reshape(B, L, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
-        k = self.k_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
-        v = self.v_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
-        q = mx.fast.rope(q, self.head_dim, traditional=False, base=self.rope_theta, scale=1.0, offset=0)
-        k = mx.fast.rope(k, self.head_dim, traditional=False, base=self.rope_theta, scale=1.0, offset=0)
-        out = mx.fast.scaled_dot_product_attention(q, k, v, mask=mask.astype(q.dtype) if mask is not None else None, scale=self.scale)
-        return self.o_proj(out.transpose(0, 2, 1, 3).reshape(B, L, -1))
-
-class MLP(nn.Module):
-    def __init__(self, dim, hidden):
-        super().__init__()
-        self.gate_proj = nn.Linear(dim, hidden, bias=False)
-        self.down_proj = nn.Linear(hidden, dim, bias=False)
-        self.up_proj = nn.Linear(dim, hidden, bias=False)
-    def __call__(self, x):
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
-
-class TransformerBlock(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        self.self_attn = Attention(args)
-        self.mlp = MLP(args.hidden_size, args.intermediate_size)
-        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-    def __call__(self, x, mask=None):
-        h = x + self.self_attn(self.input_layernorm(x), mask)
-        return h + self.mlp(self.post_attention_layernorm(h))
-
-class Qwen2Model(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
-        self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
-        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-    def __call__(self, inputs, mask=None):
-        h = self.embed_tokens(inputs)
-        for layer in self.layers:
-            h = layer(h, mask)
-        return self.norm(h)
-
-class JinaCodeEmbeddingModel(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        args = ModelArgs(
-            hidden_size=config["hidden_size"],
-            num_hidden_layers=config["num_hidden_layers"],
-            intermediate_size=config["intermediate_size"],
-            num_attention_heads=config["num_attention_heads"],
-            rms_norm_eps=config["rms_norm_eps"],
-            vocab_size=config["vocab_size"],
-            num_key_value_heads=config["num_key_value_heads"],
-            max_position_embeddings=config["max_position_embeddings"],
-            rope_theta=config.get("rope_theta", 1000000.0),
-        )
-        self.model = Qwen2Model(args)
-        self.config = config
-
-    def __call__(self, input_ids, attention_mask=None):
-        B, L = input_ids.shape
-        causal = mx.tril(mx.ones((L, L)))
-        causal = mx.where(causal == 0, -1e4, 0.0)[None, None, :, :]
-        if attention_mask is not None:
-            pad = mx.where(attention_mask == 0, -1e4, 0.0)[:, None, None, :]
-            mask = causal + pad
-        else:
-            mask = causal
-        h = self.model(input_ids, mask)
-        if attention_mask is not None:
-            seq_lens = mx.sum(attention_mask.astype(mx.int32), axis=1) - 1
-            embs = h[mx.arange(B), seq_lens]
-        else:
-            embs = h[:, -1, :]
-        norms = mx.linalg.norm(embs, axis=1, keepdims=True)
-        return embs / norms
-
-    def encode(self, texts, tokenizer, max_length=8192, truncate_dim=None, task="nl2code", prompt_type="query"):
-        PREFIXES = {"nl2code": {"query": "Find the most relevant code snippet given the following query:\n", "passage": "Candidate code snippet:\n"}}
-        prefix = PREFIXES.get(task, {}).get(prompt_type, "")
-        if prefix:
-            texts = [prefix + t for t in texts]
-        encodings = tokenizer.encode_batch(texts)
-        ml = min(max_length, max(len(e.ids) for e in encodings))
-        iids, amask = [], []
-        for e in encodings:
-            ids = e.ids[:ml]; m = e.attention_mask[:ml]; p = ml - len(ids)
-            if p > 0: ids = ids + [0]*p; m = m + [0]*p
-            iids.append(ids); amask.append(m)
-        embs = self(mx.array(iids), mx.array(amask))
-        if truncate_dim:
-            embs = embs[:, :truncate_dim]
-            embs = embs / mx.linalg.norm(embs, axis=1, keepdims=True)
-        return embs