{"id":43700779,"url":"https://github.com/jina-ai/mlx-retrieval","last_synced_at":"2026-02-05T05:04:40.064Z","repository":{"id":311247972,"uuid":"1038887977","full_name":"jina-ai/mlx-retrieval","owner":"jina-ai","description":"Train embedding and reranker models for retrieval tasks on Apple Silicon with MLX","archived":false,"fork":false,"pushed_at":"2025-09-18T06:16:34.000Z","size":239,"stargazers_count":161,"open_issues_count":0,"forks_count":8,"subscribers_count":0,"default_branch":"master","last_synced_at":"2025-10-20T07:54:10.623Z","etag":null,"topics":["apple-silicon","embeddings","mlx","mteb","reranker"],"latest_commit_sha":null,"homepage":"https://jina.ai","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/jina-ai.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-08-16T02:41:31.000Z","updated_at":"2025-10-15T02:59:32.000Z","dependencies_parsed_at":"2025-08-23T03:40:35.492Z","dependency_job_id":"f72fbdef-c43c-4a94-8867-3ae78574d5e2","html_url":"https://github.com/jina-ai/mlx-retrieval","commit_stats":null,"previous_names":["jina-ai/mlx-retrieval"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/jina-ai/mlx-retrieval","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jina-ai%2Fmlx-retrieval","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jina-ai%2Fmlx-retrieval/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jina-ai%2Fmlx-retrieval/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jina-ai%2Fmlx-retrieval/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/jina-ai","download_url":"https://codeload.github.com/jina-ai/mlx-retrieval/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/jina-ai%2Fmlx-retrieval/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":29113190,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-02-05T03:44:17.043Z","status":"ssl_error","status_checked_at":"2026-02-05T03:44:12.077Z","response_time":65,"last_error":"SSL_connect returned=1 errno=0 peeraddr=140.82.121.5:443 state=error: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["apple-silicon","embeddings","mlx","mteb","reranker"],"created_at":"2026-02-05T05:04:39.410Z","updated_at":"2026-02-05T05:04:40.059Z","avatar_url":"https://github.com/jina-ai.png","language":"Python","funding_links":[],"categories":["Training, Fine-tuning \u0026 Optimization"],"sub_categories":[],"readme":"# mlx-retrieval\n\nTrain embedding and reranker models for retrieval tasks on Apple Silicon with MLX. Features:\n\n- Full/partial LoRA training with MLX only\n- InfoNCE, NT-Xent loss with hard negative mining\n- Gradient accumulation for large batch sizes\n- MLX Data for efficient data loading\n- MTEB integration for evaluation, W\u0026B integration for logging\n\nOn M3 Ultra 512GB (80 GPU cores), training speed of `gemma-3-270m` is around 4000-5000 tokens/sec using an effective batch size of 256 and 16 gradient accumulation steps.\n\n## File Structure\n\n- `train.py` - The LoRA training script\n- `distill.py` - The post-training distillation script\n- `eval.py` - MTEB evaluation used during training; can also be used standalone\n- `embed.py` - Helper functions for generating embeddings\n- `loss.py` - Loss functions\n- `data_loader.py` - Efficient data loader that streams training data from local JSONL files or Elasticsearch indices\n- `train-data.jsonl` - Sample training data in JSONL format. This contains synthetic data for testing purposes.\n- `gemma-3-270m-mlx` - The [`gemma-3-270m`](https://huggingface.co/google/gemma-3-270m) model converted to MLX format. You can also convert it yourself using `mlx_lm.convert --hf-path unsloth/gemma-3-270m --mlx-path gemma-3-270m-mlx`. Note that `gemma-3-270m` is licensed under https://ai.google.dev/gemma/terms \n\n## Getting Started\n\nTo use the pre-converted MLX `gemma-3-270m` example model:\n\n```bash\ngit lfs install\ngit clone https://github.com/jina-ai/mlx-retrieval.git\n```\n\nTo convert the model yourself instead, skip the LFS clone and convert manually:\n\n```bash\nGIT_LFS_SKIP_SMUDGE=1 git clone https://github.com/jina-ai/mlx-retrieval.git\n```\n\n### Setup Environment\n\n```bash\n# Install uv package manager\npip install pipx\npipx install uv\n\n# Create and activate virtual environment with Python 3.12\ncd mlx-retrieval\nuv venv -p 3.12\nsource .venv/bin/activate\n\n# Install requirements\nuv pip install -r requirements.txt\n\n# (Optional) Convert the Gemma 3 270M model to MLX format if you didn't use the LFS clone\nmlx_lm.convert --hf-path unsloth/gemma-3-270m --mlx-path gemma-3-270m-mlx\n```\n\n### Training \u0026 Evaluation\n\nStart training with the following command:\n\n```bash\npython train.py \\\n    --model gemma-3-270m-mlx \\\n    --batch-size 256 \\\n    --gradient-accumulation-steps 16 \\\n    --steps 2000 \\\n    --eval-steps 100 \\\n    --save-steps 500 \\\n    --eval-tasks NanoMSMARCORetrieval \\\n    --skip-eval-init \\\n    --wandb\n```\n\nThis adds full LoRA to the model and fine-tunes it with an effective batch size of 256 over 2000 steps. Every 100 steps, the model is evaluated on the NanoMSMARCORetrieval task. The adapter is saved to the `./adapters` directory every 500 steps.\n\n![Training Progress](.readme/train.jpg)\n\nThe screenshot shows varying training tokens/sec because data is streamed from a remote Elasticsearch index. Network latency affects training speed. Using local JSONL data should provide stable training speed of 4000-5000 tokens/sec on M3 Ultra 512GB.\n\n## Distillation\n\nSee [distill-readme.md](distill-readme.md) for more details.\n\n## Technical Details\n\n- This project is primarily for educational purposes and implements common practices for training effective embedding and reranker models. While currently tested on the `gemma-3-270m` model, it should work with other models as well.\n- Unlike jina-embeddings-v3/v4, this implementation doesn't use multi-LoRA for different tasks. It implements a single LoRA configuration specifically for retrieval tasks. In v3/v4, those task-specific LoRAs are trained with different configurations and loss functions.\n- Similar to jina-embeddings-v3/v4, queries and documents are marked with \"prompt\" tokens. Queries use the format `\u003cbos\u003e\u003cunused0\u003e{text}\u003ceos\u003e`, while documents use `\u003cbos\u003e\u003cunused1\u003e{text}\u003ceos\u003e`. During embedding generation, `\u003cbos\u003e`, `\u003ceos\u003e`, and `\u003cpad\u003e` tokens are masked, but `\u003cunused0\u003e` and `\u003cunused1\u003e` tokens are preserved. The final embeddings are generated using mean pooling.\n\n## License\n\nThis project is copyright (c) 2025 Jina AI GmbH and licensed under the Apache License 2.0. The example MLX model files such as `gemma-3-270m-mlx` are third-party assets and are not covered by this project's license. Please refer to the respective model licenses for usage terms and conditions.","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fjina-ai%2Fmlx-retrieval","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fjina-ai%2Fmlx-retrieval","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fjina-ai%2Fmlx-retrieval/lists"}