Upload folder using huggingface_hub

README.md

@@ -0,0 +1,152 @@
+---
+license: apache-2.0
+library_name: mlx
+tags:
+  - mlx
+  - memory-augmented
+  - code-generation
+  - retrieval-augmented
+  - python
+  - code-search
+pipeline_tag: text-generation
+---
+
+# MALM-165M: Memory-Augmented Language Model
+
+A 165M parameter Memory-Augmented Language Model (MALM) for semantic code search, trained on CodeParrot.
+
+## Quick Start
+
+```bash
+# Install dependencies
+pip install mlx huggingface_hub numpy
+
+# Download model
+huggingface-cli download mlx-community/malm-165m --local-dir ./malm-165m
+
+# Run semantic search
+python malm-165m/inference.py --query "function that sorts a list"
+```
+
+**Example output:**
+```
+Query: function that sorts a list
+------------------------------------------------------------
+
+1. array_sort (score: 0.9526)
+   Signature: array_sort(col)
+   Docstring: Collection function: sorts the input array in ascending order...
+
+2. sort_array (score: 0.7707)
+   Signature: sort_array(col, asc)
+   Docstring: Collection function: sorts the input array in ascending or descending order...
+```
+
+## Python API
+
+```python
+from huggingface_hub import snapshot_download
+from pathlib import Path
+import sys
+
+# Download and import
+model_path = snapshot_download("mlx-community/malm-165m")
+sys.path.insert(0, model_path)
+
+from inference import load_model, search_functions
+
+# Load model
+model, tokenizer, functions, config = load_model(Path(model_path))
+print(f"Loaded {len(functions)} functions")
+
+# Search
+results = search_functions(
+    model, tokenizer, functions,
+    query="connect to database",
+    top_k=5
+)
+
+for name, signature, docstring, score in results:
+    print(f"{name}: {score:.4f}")
+```
+
+## Model Description
+
+MALM combines a transformer with learned memory retrieval for semantic code search:
+
+1. **Query encoder** - Encodes natural language queries into embeddings
+2. **Value encoder** - Encodes function signatures/docstrings
+3. **Retrieval** - Attention-based lookup from query to memory
+4. **Memory bank** - 2000 Python functions from CodeParrot
+
+### Why not mlx-lm?
+
+MALM uses a **memory-augmented** architecture different from standard LLMs:
+- Separate query and value encoders for retrieval
+- Requires a memory bank of functions
+- Inference is retrieval-based, not autoregressive generation
+
+This architecture doesn't fit `mlx-lm generate`, so we provide a custom inference script.
+
+## Architecture
+
+| Component | Parameters |
+|-----------|------------|
+| Embedding | 11.1M |
+| Position Embedding | 0.1M |
+| Query Encoder (4 layers) | 28.4M |
+| Value Encoder (4 layers) | 28.4M |
+| Decoder (12 layers) | 85.1M |
+| Output Projection | 11.1M |
+| **Total** | **~165M** |
+
+### Configuration
+
+```json
+{
+  "vocab_size": 14407,
+  "d_model": 768,
+  "n_heads": 12,
+  "n_layers": 12,
+  "n_query_layers": 4,
+  "max_seq_len": 128,
+  "num_parameters": 165123656,
+  "num_functions": 2000
+}
+```
+
+## Files
+
+| File | Description |
+|------|-------------|
+| `model.npz` | Model weights (MLX-compatible NumPy format) |
+| `config.json` | Model configuration |
+| `tokenizer.json` | Tokenizer vocabulary |
+| `functions.json` | Memory bank of 2000 Python functions |
+| `inference.py` | Standalone inference script |
+
+## Training
+
+Trained on CodeParrot with a focus on Python function retrieval:
+- Encodes natural language queries into embedding space
+- Learns semantic similarity between queries and function signatures
+- Uses attention-based retrieval over a memory bank
+
+## Related Work
+
+Part of the [HashHop](https://github.com/codelion/hash-hop) project exploring long-context evaluation and memory-augmented architectures.
+
+## Citation
+
+```bibtex
+@misc{malm2025,
+  title={MALM: Memory-Augmented Language Model},
+  author={HashHop Contributors},
+  year={2025},
+  url={https://github.com/codelion/hash-hop}
+}
+```
+
+## License
+
+Apache 2.0

config.json

@@ -0,0 +1,10 @@
+{
+  "vocab_size": 14407,
+  "d_model": 768,
+  "n_heads": 12,
+  "n_layers": 12,
+  "n_query_layers": 4,
+  "max_seq_len": 128,
+  "num_parameters": 165123656,
+  "num_functions": 2000
+}

inference.py

@@ -0,0 +1,286 @@
+#!/usr/bin/env python3
+"""MALM Inference Script - Run directly from Hugging Face.
+
+Usage:
+    # Install dependencies
+    pip install mlx huggingface_hub
+
+    # Download and run
+    huggingface-cli download mlx-community/malm-165m --local-dir ./malm-165m
+    python malm-165m/inference.py --query "function that sorts a list"
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+import json
+import argparse
+from pathlib import Path
+from typing import List, Dict, Tuple
+import re
+
+
+class MALM(nn.Module):
+    """Memory-Augmented Language Model."""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int = 768,
+        n_heads: int = 12,
+        n_layers: int = 12,
+        n_query_layers: int = 4,
+        max_seq_len: int = 128,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.n_query_layers = n_query_layers
+        self.max_seq_len = max_seq_len
+
+        # Embeddings
+        self.embed = nn.Embedding(vocab_size, d_model)
+        self.pos_embed = nn.Embedding(max_seq_len, d_model)
+        self.embed_dropout = nn.Dropout(dropout)
+
+        # Query encoder
+        self.query_layers = [
+            nn.TransformerEncoderLayer(d_model, n_heads, d_model * 4)
+            for _ in range(n_query_layers)
+        ]
+        self.query_ln = nn.LayerNorm(d_model)
+        self.query_proj = nn.Linear(d_model, d_model)
+
+        # Value encoder
+        self.value_layers = [
+            nn.TransformerEncoderLayer(d_model, n_heads, d_model * 4)
+            for _ in range(n_query_layers)
+        ]
+        self.value_ln = nn.LayerNorm(d_model)
+        self.value_proj = nn.Linear(d_model, d_model)
+
+        # Decoder layers
+        self.decoder_layers = [
+            nn.TransformerEncoderLayer(d_model, n_heads, d_model * 4)
+            for _ in range(n_layers)
+        ]
+        self.decoder_ln = nn.LayerNorm(d_model)
+
+        # Output
+        self.output = nn.Linear(d_model, vocab_size)
+
+        # Temperature for retrieval
+        self.log_temp = mx.array([0.0])
+
+    def encode_query(self, query_ids: mx.array) -> mx.array:
+        """Encode query to single embedding."""
+        B, L = query_ids.shape
+
+        h = self.embed(query_ids)
+        pos = mx.arange(min(L, self.max_seq_len))
+        h = h + self.pos_embed(pos)
+        h = self.embed_dropout(h)
+
+        for layer in self.query_layers:
+            h = layer(h, None)
+
+        h = self.query_ln(h)
+
+        mask = (query_ids != 0).astype(mx.float32)[:, :, None]
+        h = h * mask
+        query_emb = mx.sum(h, axis=1) / (mx.sum(mask, axis=1) + 1e-8)
+
+        return self.query_proj(query_emb)
+
+    def encode_value(self, value_ids: mx.array) -> mx.array:
+        """Encode value to single embedding."""
+        B, L = value_ids.shape
+
+        h = self.embed(value_ids)
+        pos = mx.arange(min(L, self.max_seq_len))
+        h = h + self.pos_embed(pos)
+
+        for layer in self.value_layers:
+            h = layer(h, None)
+
+        h = self.value_ln(h)
+
+        mask = (value_ids != 0).astype(mx.float32)[:, :, None]
+        h = h * mask
+        val_emb = mx.sum(h, axis=1) / (mx.sum(mask, axis=1) + 1e-8)
+
+        return self.value_proj(val_emb)
+
+    def retrieve(
+        self,
+        query_emb: mx.array,
+        key_emb: mx.array,
+        val_emb: mx.array,
+    ) -> Tuple[mx.array, mx.array, mx.array]:
+        """Retrieve from memory."""
+        scale = self.d_model ** -0.5
+        temp = mx.exp(self.log_temp) + 0.1
+
+        scores = (query_emb @ key_emb.T) * scale / temp
+        attn = mx.softmax(scores, axis=-1)
+        retrieved = attn @ val_emb
+
+        return retrieved, attn, scores
+
+
+class Tokenizer:
+    """Simple tokenizer for MALM."""
+
+    def __init__(self, tokenizer_dict: Dict):
+        self.token_to_id = tokenizer_dict.get("token_to_id", {})
+        self.id_to_token = {int(v): k for k, v in self.token_to_id.items()}
+        self.special = {"<PAD>": 0, "<UNK>": 1, "<BOS>": 2, "<EOS>": 3}
+
+    def encode(self, text: str) -> List[int]:
+        """Tokenize text."""
+        tokens = re.findall(r"[a-zA-Z_][a-zA-Z0-9_]*|[0-9]+|[^\s]", text.lower())
+        return [self.token_to_id.get(t, self.special.get("<UNK>", 1)) for t in tokens]
+
+    def decode(self, ids: List[int]) -> str:
+        """Decode token IDs to text."""
+        tokens = [self.id_to_token.get(i, "<UNK>") for i in ids]
+        return " ".join(tokens)
+
+
+def load_model(model_dir: Path):
+    """Load MALM model from directory."""
+    import mlx.utils as mlx_utils
+
+    # Load config
+    with open(model_dir / "config.json") as f:
+        config = json.load(f)
+
+    # Create model
+    model = MALM(
+        vocab_size=config["vocab_size"],
+        d_model=config["d_model"],
+        n_heads=config["n_heads"],
+        n_layers=config["n_layers"],
+        n_query_layers=config["n_query_layers"],
+        max_seq_len=config["max_seq_len"],
+    )
+
+    # Load weights and convert to mlx arrays
+    weights = dict(np.load(model_dir / "model.npz"))
+    weights = {k: mx.array(v) for k, v in weights.items()}
+
+    # Unflatten and load
+    params = mlx_utils.tree_unflatten(list(weights.items()))
+    model.update(params)
+    mx.eval(model.parameters())
+
+    # Load tokenizer
+    with open(model_dir / "tokenizer.json") as f:
+        tokenizer_dict = json.load(f)
+    tokenizer = Tokenizer(tokenizer_dict)
+
+    # Load functions
+    with open(model_dir / "functions.json") as f:
+        functions = json.load(f)
+
+    return model, tokenizer, functions, config
+
+
+def search_functions(
+    model: MALM,
+    tokenizer: Tokenizer,
+    functions: List[Dict],
+    query: str,
+    top_k: int = 5,
+) -> List[Tuple[str, str, float]]:
+    """Search for functions matching a query.
+
+    Uses the function name as key and signature+docstring as value for retrieval.
+    """
+    # Encode query
+    query_ids = tokenizer.encode(query)
+    if not query_ids:
+        query_ids = [1]  # <UNK>
+    query_ids = mx.array([query_ids])
+
+    # Encode all function keys and values
+    key_tokens = []
+    value_tokens = []
+    max_val_len = 64
+
+    for func in functions:
+        name = func["name"]
+        # Use signature + docstring as the "value" to search over
+        sig = func.get("signature", name)
+        doc = func.get("docstring", "")
+        value_text = f"{sig} {doc}"
+
+        key_id = tokenizer.token_to_id.get(name.lower(), 1)
+        key_tokens.append(key_id)
+
+        val_ids = tokenizer.encode(value_text)[:max_val_len]
+        val_ids = val_ids + [0] * (max_val_len - len(val_ids))
+        value_tokens.append(val_ids)
+
+    key_tokens = mx.array(key_tokens)
+    value_tokens = mx.array(value_tokens)
+
+    # Encode memory
+    key_emb = model.embed(key_tokens)
+    val_emb = model.encode_value(value_tokens)
+
+    # Get query embedding and compute similarity
+    query_emb = model.encode_query(query_ids)
+    _, attn, scores = model.retrieve(query_emb, key_emb, val_emb)
+    mx.eval(scores)
+
+    # Get top-k
+    scores_np = np.array(scores[0])
+    top_indices = np.argsort(scores_np)[::-1][:top_k]
+
+    results = []
+    for idx in top_indices:
+        func = functions[idx]
+        score = float(scores_np[idx])
+        sig = func.get("signature", func["name"])
+        doc = func.get("docstring", "")
+        results.append((func["name"], sig, doc, score))
+
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MALM Inference - Semantic Code Search")
+    parser.add_argument("--query", type=str, required=True, help="Natural language query")
+    parser.add_argument("--top-k", type=int, default=5, help="Number of results")
+    parser.add_argument("--model-dir", type=str, default=None, help="Model directory")
+    args = parser.parse_args()
+
+    # Determine model directory
+    if args.model_dir:
+        model_dir = Path(args.model_dir)
+    else:
+        model_dir = Path(__file__).parent
+
+    print(f"Loading model from {model_dir}...")
+    model, tokenizer, functions, config = load_model(model_dir)
+    print(f"Loaded {len(functions)} functions, {config['num_parameters']:,} parameters")
+
+    # Search
+    print(f"\nQuery: {args.query}")
+    print("-" * 60)
+
+    results = search_functions(model, tokenizer, functions, args.query, args.top_k)
+
+    for i, (name, signature, docstring, score) in enumerate(results, 1):
+        print(f"\n{i}. {name} (score: {score:.4f})")
+        print(f"   Signature: {signature}")
+        if docstring:
+            print(f"   Docstring: {docstring[:100]}{'...' if len(docstring) > 100 else ''}")
+
+
+if __name__ == "__main__":
+    main()

model.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f7c9ba2754706b3e7c9c0b00b53226a6bafdcb2aaab5b47a46bdf5d21e4f14f
+size 660569734