langchain/docs/examples/prompts/custom_llm.ipynb
2022-11-26 06:07:02 -08:00

154 lines
3.4 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"id": "9e9b7651",
"metadata": {},
"source": [
"# Custom LLM\n",
"\n",
"This notebook goes over how to create a custom LLM wrapper, in case you want to use your own LLM or a different wrapper than one that is supported in LangChain.\n",
"\n",
"There is only one required thing that a custom LLM needs to implement:\n",
"\n",
"1. A `__call__` method that takes in a string, some optional stop words, and returns a string\n",
"\n",
"There is a second optional thing it can implement:\n",
"\n",
"1. An `_identifying_params` property that is used to help with printing of this class. Should return a dictionary.\n",
"\n",
"Let's implement a very simple custom LLM that just returns the first N characters of the input."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "a65696a0",
"metadata": {},
"outputs": [],
"source": [
"from langchain.llms.base import LLM\n",
"from typing import Optional, List, Mapping, Any"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "d5ceff02",
"metadata": {},
"outputs": [],
"source": [
"class CustomLLM(LLM):\n",
" \n",
" def __init__(self, n: int):\n",
" self.n = n\n",
" \n",
" def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str:\n",
" if stop is not None:\n",
" raise ValueError(\"stop kwargs are not permitted.\")\n",
" return prompt[:self.n]\n",
" \n",
" @property\n",
" def _identifying_params(self) -> Mapping[str, Any]:\n",
" \"\"\"Get the identifying parameters.\"\"\"\n",
" return {\"n\": self.n}"
]
},
{
"cell_type": "markdown",
"id": "714dede0",
"metadata": {},
"source": [
"We can now use this as an any other LLM."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "10e5ece6",
"metadata": {},
"outputs": [],
"source": [
"llm = CustomLLM(n=10)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "8cd49199",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'This is a '"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"llm(\"This is a foobar thing\")"
]
},
{
"cell_type": "markdown",
"id": "bbfebea1",
"metadata": {},
"source": [
"We can also print the LLM and see its custom print."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "9c33fa19",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[1mCustomLLM\u001b[0m\n",
"Params: {'n': 10}\n"
]
}
],
"source": [
"print(llm)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6dac3f47",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.6"
}
},
"nbformat": 4,
"nbformat_minor": 5
}