Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add mlx support #1089

Merged
merged 16 commits into from
Jan 10, 2025
Merged

add mlx support #1089

merged 16 commits into from
Jan 10, 2025

Conversation

davidberenstein1957
Copy link
Member

@davidberenstein1957 davidberenstein1957 commented Dec 30, 2024

Closes #995

Use it individually

from distilabel.models.llms import MlxLLM

llm = MlxLLM(model="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit")

llm.load()

# Call the model
output = llm.generate_outputs(inputs=[[{"role": "user", "content": "Hello world!"}]])

Use it with magpie

from distilabel.models.llms.mlx import MlxLLM
from distilabel.steps.tasks import Magpie

model = MlxLLM(
    path_or_hf_repo="mlx-community/Meta-Llama-3.1-8B-Instruct-8bit",
    use_magpie_template=True,
    magpie_pre_query_template="llama3",
)
task = Magpie(llm=model)
task.load()

It is relatively easy to spin up an mlx server, but no public Python API clients are available except for LangChain https://python.langchain.com/docs/integrations/chat/mlx/. Currently, the OpenAIAPI does not align with the payloads from both the chat and the text generation API.

Copy link

Documentation for this PR has been built. You can view it at: https://distilabel.argilla.io/pr-1089/

Copy link

codspeed-hq bot commented Dec 30, 2024

CodSpeed Performance Report

Merging #1089 will not alter performance

Comparing feat/add-support-mlx (a65a268) with develop (99c2448)

Summary

✅ 1 untouched benchmarks

@davidberenstein1957 davidberenstein1957 marked this pull request as ready for review January 8, 2025 15:48
@davidberenstein1957 davidberenstein1957 linked an issue Jan 8, 2025 that may be closed by this pull request
Copy link
Member

@gabrielmbmb gabrielmbmb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's also add the import of MlxLLM to src/distilabel/llms.py to avoid confusion until we deprecate it in 1.7.0

src/distilabel/models/llms/mlx.py Outdated Show resolved Hide resolved
src/distilabel/models/llms/mlx.py Outdated Show resolved Hide resolved
davidberenstein1957 and others added 7 commits January 10, 2025 11:38
- Introduced MlxLLM class in mlx.py, integrating it into the llms module.
- Updated output preparation logic to include token computation and logprobs in utils.py.
- Modified __init__.py to export MlxLLM.
- Enhanced type annotations for tokenizer_config and model_config in MlxLLM.
- Imported LlamaCppEmbeddings in models/__init__.py.
- Added LlamaCppEmbeddings to the __all__ exports in both models/__init__.py and embeddings/__init__.py.
- Removed duplicate entry of LlamaCppEmbeddings from embeddings/__init__.py exports.
@davidberenstein1957 davidberenstein1957 merged commit 2c893c1 into develop Jan 10, 2025
8 checks passed
@davidberenstein1957 davidberenstein1957 deleted the feat/add-support-mlx branch January 10, 2025 11:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[FEATURE] mlx-lm integration
2 participants