File size: 3,778 Bytes
cfd3735
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9e9b7651",
   "metadata": {},
   "source": [
    "# How to write a custom LLM wrapper\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": 6,
   "id": "a65696a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, List, Mapping, Optional\n",
    "\n",
    "from langchain.callbacks.manager import CallbackManagerForLLMRun\n",
    "from langchain.llms.base import LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d5ceff02",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomLLM(LLM):\n",
    "    \n",
    "    n: int\n",
    "        \n",
    "    @property\n",
    "    def _llm_type(self) -> str:\n",
    "        return \"custom\"\n",
    "    \n",
    "    def _call(\n",
    "        self,\n",
    "        prompt: str,\n",
    "        stop: Optional[List[str]] = None,\n",
    "        run_manager: Optional[CallbackManagerForLLMRun] = None,\n",
    "    ) -> 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": 8,
   "id": "10e5ece6",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = CustomLLM(n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8cd49199",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'This is a '"
      ]
     },
     "execution_count": 9,
     "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": 10,
   "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}