MLX
This notebook shows how to get started usingMLX
LLM's as chat models.
In particular, we will:
- Utilize theMLXPipeline,
- Utilize the
ChatMLX
class to enable any of these LLMs to interface with LangChain'sChat Messages abstraction. - Demonstrate how to use an open-source LLM to power an
ChatAgent
pipeline
%pip install--upgrade--quiet mlx-lm transformers huggingface_hub
1. Instantiate an LLM
There are three LLM options to choose from.
from langchain_community.llms.mlx_pipelineimport MLXPipeline
llm= MLXPipeline.from_model_id(
"mlx-community/quantized-gemma-2b-it",
pipeline_kwargs={"max_tokens":10,"temp":0.1},
)
API Reference:MLXPipeline
2. Instantiate theChatMLX
to apply chat templates
Instantiate the chat model and some messages to pass.
from langchain_community.chat_models.mlximport ChatMLX
from langchain_core.messagesimport HumanMessage
messages=[
HumanMessage(
content="What happens when an unstoppable force meets an immovable object?"
),
]
chat_model= ChatMLX(llm=llm)
API Reference:ChatMLX |HumanMessage
Inspect how the chat messages are formatted for the LLM call.
chat_model._to_chat_prompt(messages)
Call the model.
res= chat_model.invoke(messages)
print(res.content)
3. Take it for a spin as an agent!
Here we'll test outgemma-2b-it
as a zero-shotReAct
Agent. The example below is taken fromhere.
Note: To run this section, you'll need to have aSerpAPI Token saved as an environment variable:
SERPAPI_API_KEY
from langchainimport hub
from langchain.agentsimport AgentExecutor, load_tools
from langchain.agents.format_scratchpadimport format_log_to_str
from langchain.agents.output_parsersimport(
ReActJsonSingleInputOutputParser,
)
from langchain.tools.renderimport render_text_description
from langchain_community.utilitiesimport SerpAPIWrapper
API Reference:hub |AgentExecutor |load_tools |format_log_to_str |ReActJsonSingleInputOutputParser |render_text_description |SerpAPIWrapper
Configure the agent with areact-json
style prompt and access to a search engine and calculator.
# setup tools
tools= load_tools(["serpapi","llm-math"], llm=llm)
# setup ReAct style prompt
# Based on 'hwchase17/react' prompt modification, cause mlx does not support the `System` role
human_prompt="""
Answer the following questions as best you can. You have access to the following tools:
{tools}
The way you use the tools is by specifying a json blob.
Specifically, this json should have a `action` key (with the name of the tool to use) and a `action_input` key (with the input to the tool going here).
The only values that should be in the "action" field are: {tool_names}
The $JSON_BLOB should only contain a SINGLE action, do NOT return a list of multiple actions. Here is an example of a valid $JSON_BLOB:
\`\`\`
{{
"action": $TOOL_NAME,
"action_input": $INPUT
}}
\`\`\`
ALWAYS use the following format:
Question: the input question you must answer
Thought: you should always think about what to do
Action:
\`\`\`
$JSON_BLOB
\`\`\`
Observation: the result of the action
... (this Thought/Action/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question
Begin! Reminder to always use the exact characters `Final Answer` when responding.
{input}
{agent_scratchpad}
"""
prompt= human_prompt.partial(
tools=render_text_description(tools),
tool_names=", ".join([t.namefor tin tools]),
)
# define the agent
chat_model_with_stop= chat_model.bind(stop=["\nObservation"])
agent=(
{
"input":lambda x: x["input"],
"agent_scratchpad":lambda x: format_log_to_str(x["intermediate_steps"]),
}
| prompt
| chat_model_with_stop
| ReActJsonSingleInputOutputParser()
)
# instantiate AgentExecutor
agent_executor= AgentExecutor(agent=agent, tools=tools, verbose=True)
agent_executor.invoke(
{
"input":"Who is Leo DiCaprio's girlfriend? What is her current age raised to the 0.43 power?"
}
)
Related
- Chat modelconceptual guide
- Chat modelhow-to guides