|
1 | | -# PGML Discord Bot Example |
| 1 | +# Discord Bot using pgml Python SDK, Lanchain, Instructor-xl, and Falcon 7B |
2 | 2 |
|
3 | | -Temporary readme while I work on the tutorial. |
| 3 | +In this tutorial, we will build a Discord bot that can use markdown files to help answer user inquiries. We will ingest the files, convert their contents into vector embeddings, and save them to Postgres. After indexing the data, the bot will query the collection to retrieve the documents that are most likely to answer the user's question. Then, we will use a simple SQL query utilizing PostgresML to retrieve a completion from the open source Falcon 7b text generation model. Finally, we will return this completion to the user in the Discord channel. We will be using the pgml SDK to simplify the process. |
4 | 4 |
|
5 | | -There are three main files: |
| 5 | +In this project, we will be working with three files: |
6 | 6 |
|
7 | | -1. `./ingest_data.ipynb` - Jupyter notebook you can run to ingest your data and index |
8 | | -2. `./discord_bot.py` - The discord bot class |
9 | | -3. `./start_bot.py` - File that starts the bot |
| 7 | +1. `./ingest.ipynb` - Jupyter notebook you can run to ingest your data into the bot. |
| 8 | +2. `./bot.py` - The bot itself |
| 9 | +3. `./start.py` - File that starts the bot. |
| 10 | + |
| 11 | +## Step 1: Create a Bot and Set Up Your Environment |
| 12 | + |
| 13 | +First, we will ensure that we have all the necessary environment variables set. |
| 14 | + |
| 15 | +Make a copy of the `.env.template` file and name it `.env`, and ensure it is located in the root directory. |
| 16 | + |
| 17 | +Now, we will go through each of the variables and set them to the appropriate values. |
| 18 | + |
| 19 | +To create a Discord bot, you will need to create a Discord bot account and get a token. You can follow the tutorial on how to do that [here](https://discordpy.readthedocs.io/en/stable/discord.html). After going through this tutorial, you should have a bot created and added to your server. You should also have a token for your bot. Set this token to the variable `DISCORD_TOKEN` in your .env file. |
| 20 | + |
| 21 | +Next, set the name of the Discord channel you would like the bot to listen to. Set this to the variable `DISCORD_CHANNEL` in your .env file. |
| 22 | + |
| 23 | +We will be using the pgml Python SDK to create, store, and query our vectors. So, if you don't already have an account there, you can create one here: https://postgresml.org/. You can select the free serverless option and will be given a connection string. Set this connection string to the variable `PGML_CONNECTION_STR` in your .env file. |
| 24 | + |
| 25 | +Next, you will want to add the markdown files you would like to use into the `./content` folder. Set the path to this folder to the variable `CONTENT_PATH` in your .env file. |
| 26 | + |
| 27 | +Now that our project is set up, we can start working on our bot. |
| 28 | + |
| 29 | +## Step 2: Ingest Your Data |
| 30 | + |
| 31 | +Now we will ingest our markdown data into the bot. |
| 32 | + |
| 33 | +Open and run the cells in the `./ingest.ipynb` notebook. If you have set all of your environment variables correctly, and put your markdown files into your markdown folder correctly, you should be able to run the notebook without any errors, and your bot will now have access to your data. |
| 34 | + |
| 35 | +Let's take a look at what is happening in the notebook. |
| 36 | + |
| 37 | +1. We load in the markdown files from the path we passed in, using Lanchain's document loader. |
| 38 | +2. We convert this array of documents to an array of dictionaries in the format expected by the PGML SDK. |
| 39 | + |
| 40 | +``` |
| 41 | +
|
| 42 | +docs = [{text: 'foo'}, {text: 'bar'}, ...] |
| 43 | +
|
| 44 | +``` |
| 45 | + |
| 46 | +1. We create a PGML collection upserting those documents into a PostgreML Collection. |
| 47 | + |
| 48 | +``` |
| 49 | +
|
| 50 | +collection = pgml.create_or_get_collection(collection_name) |
| 51 | +
|
| 52 | +collection.upsert_documents(docs) |
| 53 | +
|
| 54 | +``` |
| 55 | + |
| 56 | +1. We chunk those documents into smaller sizes and embed those chunks using the Instructor-XL model. |
| 57 | + |
| 58 | +``` |
| 59 | +
|
| 60 | +collection.generate_chunks() |
| 61 | +
|
| 62 | +collection.generate_embeddings() |
| 63 | +
|
| 64 | +``` |
| 65 | + |
| 66 | +Now that our data is properly indexed, we can start our bot server to handle incoming requests, using the data we just ingested to help answer questions. |
| 67 | + |
| 68 | +## Run Your Bot |
| 69 | + |
| 70 | +For our bot server, we are using the popular library [discord.py](https://discordpy.readthedocs.io/en/stable/). |
| 71 | + |
| 72 | +To start the bot server, you can run the following command in your terminal: |
| 73 | + |
| 74 | +``` |
| 75 | +
|
| 76 | +python start.py |
| 77 | +
|
| 78 | +``` |
| 79 | + |
| 80 | +If everything was set up correctly in earlier steps, your bot should be fully functional. |
| 81 | + |
| 82 | +But since it's good to know how things are working, let's take a look at the code. |
| 83 | + |
| 84 | +In the `start.py` file, you will see the following code: |
| 85 | + |
| 86 | +``` |
| 87 | +
|
| 88 | +# get environment variables |
| 89 | +
|
| 90 | +pg_connection_string = os.getenv("PGML_CONNECTION_STR") |
| 91 | +
|
| 92 | +# ... |
| 93 | +
|
| 94 | +## initialize bot |
| 95 | +
|
| 96 | +pgml_bot = Bot(conninfo=pg_connection_string) |
| 97 | +
|
| 98 | +## start discord bot |
| 99 | +
|
| 100 | +pgml_bot.start(collection_name, discord_token) |
| 101 | +
|
| 102 | +``` |
| 103 | + |
| 104 | +This code will initialize the bot class with your PostgreSQL connection string and then start the Discord bot with the collection name, from which you previously saved your data in the previous step, and Discord token. |
| 105 | + |
| 106 | +If we look in the `.start` method, we will see that we execute `.run` on the `discord_client` which has been initialized. |
| 107 | + |
| 108 | +We also declared the `on_message` function that is called when a message is sent in the Discord server. |
| 109 | + |
| 110 | +When a message is handled by this `on_message` function, we do a few things: |
| 111 | + |
| 112 | +1. Using the PGML SDK, we run: |
| 113 | + |
| 114 | +``` |
| 115 | +
|
| 116 | +collection.vector_search( |
| 117 | +
|
| 118 | +query, |
| 119 | +
|
| 120 | +top_k=3, |
| 121 | +
|
| 122 | +model_id=2, |
| 123 | +
|
| 124 | +splitter_id=2, |
| 125 | +
|
| 126 | +query_parameters={"instruction": "Represent the question for retrieving supporting documents: "}, |
| 127 | +
|
| 128 | +) |
| 129 | +
|
| 130 | +``` |
| 131 | + |
| 132 | +This is going to return the top 3 documents that are most similar to the user's message. |
| 133 | + |
| 134 | +1. We then concatenate the text of those documents into a single string and add it to our prompt text, which looks like: |
| 135 | + |
| 136 | +```` |
| 137 | +
|
| 138 | +Use the context, which is delimited by three back ticks, below to help answer the question. |
| 139 | +
|
| 140 | +context: ```{context}``` |
| 141 | +
|
| 142 | +{user_message} |
| 143 | +
|
| 144 | +```` |
| 145 | + |
| 146 | +1. Now that we have our prompt ready, we can make a Falcon completion. We will get this completion by executing a SQL query that uses `pgml.transform` function. |
| 147 | + |
| 148 | +``` |
| 149 | +
|
| 150 | +async def run_transform_sql(self, context, message_content): |
| 151 | +
|
| 152 | +prompt = self.prepare_prompt(context, message_content) |
| 153 | +
|
| 154 | +sql_query = """SELECT pgml.transform( |
| 155 | +
|
| 156 | +task => '{ |
| 157 | +
|
| 158 | +"model": "tiiuae/falcon-7b-instruct", |
| 159 | +
|
| 160 | +"device_map": "auto", |
| 161 | +
|
| 162 | +"torch_dtype": "bfloat16", |
| 163 | +
|
| 164 | +"trust_remote_code": true |
| 165 | +
|
| 166 | +}'::JSONB, |
| 167 | +
|
| 168 | +args => '{ |
| 169 | +
|
| 170 | +"max_new_tokens": 100 |
| 171 | +
|
| 172 | +}'::JSONB, |
| 173 | +
|
| 174 | +inputs => ARRAY[%s] |
| 175 | +
|
| 176 | +) AS result""" |
| 177 | +
|
| 178 | +sql_params = (prompt,) |
| 179 | +
|
| 180 | +return await self.run_query(sql_query, sql_params) |
| 181 | +
|
| 182 | +``` |
| 183 | + |
| 184 | +1. Now that we have the response from Falcon, we need to clean the response text up a bit before returning the bot's answer. Since the completion text includes the original prompt, we will remove that from the generated text in the `prepare_response` function. |
| 185 | +2. Finally, we will send the response back to the Discord channel. |
| 186 | + |
| 187 | +## Final Remarks |
| 188 | + |
| 189 | +At this point, you should have a functioning Discord bot that can answer questions based on the markdown files you have provided, using fully open-source tools. |
| 190 | + |
| 191 | +If you have any questions or would like to chat, you can join us on the [PostgresML Discord](https://discord.gg/DmyJP3qJ7U) |
0 commit comments