Build a Telegram chatbot with any AI model under the hood

tl;dr: we are building a text summarization bot for long articles, but I keep the code as general as possible, so consider it as a Python template for wrapping any model into a bot. Don’t want to read the post? Go to the github repo. Or read a summary of the post produced by this bot at the end of the page!

In this post, I want to show how to cross a Telegram chatbot with an AI model using Python, Hugging Face🤗 and Pyppeteer. Why? Because this scenario has unlimited applications! We’ll run everything locally, but as a bonus, you will get a Dockerfile and a Python package that you can use in any cloud provider.

We live in the era of AI, and every week there is yet another groundbreaking news about what AI has learned to do. We are super lucky that many of such models are shared publicly, supported by the community and available for anyone for free.
On the other hand, Telegram is a rock star across messengers, and among many other features, it provides a powerful Bot platform, handy SDK and, of course, a clean messaging interface. One can build rich apps and even games with the Bot SDK, but in this blog post, we will build a basic chatbot with AI functionality under the hood.

Idea: Article summarization bot

I often find a good-looking but long article on the Internet that “I’ll definitely read later”. So, I peacefully send a URL of an article to my stash of links (everyone has a place like this, right?), and, of course, most of the time, I never read it.

So, I thought it would be nice to have a chat bot where I send a URL of a page, and the bot sends me back a summary of it. Sounds neat! It would be super handy since I don’t need to spend time reading the whole text, but I’ll still get all the main points, and, in case it catches my eye, I’ll definitely read the whole story later.

Ok, so we have an idea of what this thing should do. Now, let’s break it down into pieces.

The simple architecture of the Article Summarization Telegram Bot

1. We need a convenient Python library that makes all the Telegram communication stuff. After a short research, I picked aiogram. It is well-maintained, asynchronous (and fast) and has all the necessary features
2. We need a library that scrapes the content of a URL. Moreover, it should work with the pages manipulated by JavaScript since most websites have some JS going on before showing you the content of a page. That means it’s not enough to use something simple like request or beautifulsoup. Instead, we will use the full-fledged headless web browser Pyppeteer that I love with all my heart.
3. One tricky piece is that we need a library that extracts the main content from the page. If you think about it, the page has many noisy elements (headers, footers, ads, parts of design), so it is not obvious how to find the actual content text given an arbitrary HTML page. After a quick research, I found a library trafilatura, which includes several algorithms for content extraction and settings that can be adjusted
4. The heart of the project is the Text Summarization model. I experimented with a few of the most popular PyTorch models based on BART from the Hugging Face repo: facebook/bart-large-cnn (large model) and sshleifer/distilbart-cnn-12–6 (lighter distilled model). I also made the model name to be a parameter that can be changed if necessary. One a bit tricky part here will be the handling of long articles since models usually have the maximum input size
5. Packaging, testing and deployment. Here we need to prepare a minimal reproducible code to run the bot somewhere except my local machine (for example, in Google Cloud, Azure or on your local machine). To frame the project, we will create a Dockerfile and Python package with a few unit tests

Let’s code!

Great, we have 5 independent pieces of work to do here. I will go through the code from the project repository and explain all major functions. If some code is missing in the current blog post, please refer to the source code.

Part 0: Setup the project and install dependencies

To start, you need to create a bot with a @BotFather maker and copy the token. The process is relatively straightforward: come up with a name and username for your bot and save the token somewhere.

Creating a Telegram bot with @BotFater bot maker

Then, let’s create a project and install dependencies. I’m using PyCharm editor and highly recommend it. Create a project, clone the repo and set up all necessary packages from the requirements.txt.

Then go to the settings.py and modify it accordingly to your project. There are 4 blocks to configure: log settings, Telegram bot settings, Pyppeteer settings and model settings. The minimal step here is to insert your token, which is taken from the BOT_TOKEN environment variable.

Part 1: Telegram bot essentials

All bot-related code is located in the summary_bot/bot.py module. As you can see, thanks to the aiogram library, the code is readable and simple. There are 3 blocks:

1. Bot and Dispatcher objects to communicate with the Telegram backend and receive updates for your bot from users
2. One function (dispatcher handler) welcome() is called when the bot is started. It will send a Hello message to a newcomer
3. The heart of the bot — the function get_summary() for processing the messages sent to the bot. This function runs all the logic: a naive check that the message is actually a URL, then runs the web crawler, applies the summarization model and sends back to the user a summary of the URL content with message.reply()

To create a new asynchronous handler, you need to define a function with the argument aiogram.types.Message and decorate it with @dp.message_handler().

See the implementation of all 3 blocks here or summary_bot/bot.py:

Code for the Telegram bot initialization with two handlers

Part 2,3: Web Scraping and Content Extraction

All the web scraping code can be found here: summary_bot/scraper.py.

Web Scrapping code with Pyppeteer and trafilatura libraries

On line #18, there is a main function extract_article() that is called from the bot. We create a browser, get the page with a given URL, apply the function extract(content) to extract the web content with the trafilatura package (with the favor_recall setting to increase the chance of scraping the whole text) and close the browser.

On line #53, we define the Scraper class to include there all the methods.

On line #94, there is a main internal scraping method get_response(url) which contains the standard Pyppeteer/Puppeteer code for scraping.

On line #44, we define a very nasty function kill_browser(), and please forget about this as soon as you see it. Basically, sometimes the puppeteer browser can crash or leave a zombie browser even after it is closed. If you scrape something often and create browsers, you can end up with a state where the browser is broken, or there is an army of browser zombies, and your computer will run out of memory. As a workaround for this, I severely kill the browser process in case it crashes, which can create a fork of the current process.

Part 4: Text summarization model

This is part of the project where you can inject any model from the Hugging Face🤗 repo, not necessarily for a summarization task. You can process the user’s image with some fancy Computer Vision model or apply another NLP model to the text. I picked the Summarization task as an example.

The model-related code can be seen in summary_bot/summarizer.py.

Article Summarization code with HuggingFace BART model

Let’s quickly see what is going on here.

On line #89 the main function summarize_article() is defined. That is the function that is called from the bot. These are the steps in it:

1. Create the instance of a SummaryModel, which is our custom class with handy methods. It is a singleton, so we do not create a new model on every request
2. We want to process both long and short articles, but AI models often have the maximum number of input tokens. So, with the model.prepare_batch() method, we solve it by dividing the long text into chunks of max token size (1024). Then we prepare the batch of input token ids for the model (that’s usually the input for an NLP model). The caveat here is that instead of preparing one batch of size (n_chunks, 1024) and zero padding for the last chunk, we will create a tuple of batches: (all chunks except the last one, the last one) (see line #61). It leads to better summarization results for the last chunk than with zero padding.
3. Apply the summarization model with model.get_summary(). Essentially, our bot will give us a summary of pieces of max token size .
4. Prettify the final summary a bit to produce a nice list of summary bullets

On line #18, you see the SummaryModel class definition. There are 3 main methods: prepare_batch() and get_summary() we’ve discussed, and the load(), where the model and the corresponding tokenizer are loaded. We used BartForConditionalGeneration loader that is used for a summary generation with BART models (facebook/bart-large-cnn and sshleifer/distilbart-cnn-12–6)

Part 5: Packaging, testing and deployment

One option to run the bot is to run it locally from the earlier created local python environment with the command python summary_bot/bot.py

It will start the process, and you can go to the Telegram window with your bot and send some URL to check how it works.

To make the code reproducible, it is handy to prepare a Dockerfile with all instructions and dependency installation required for the bot to run.
So, another option is to run the bot from the container. It is preferable since you don’t need to care about dependencies.

Dockerfile for deploying the Summarization Bot

On line #4, there is a tricky part related to running the headless web browser puppeteer (from Python, the wrapper is the Pyppeteer). It worked seamlessly locally, but I spent not one hour to make it finally work from Docker. It was constantly giving me the “browser closed unexpectedly” exception. The problem is that when running in Docker, puppeteer requires a non-default browser executable and some extra launching settings. The fact that the bot is running from the container is passed via the environment variable FROM_DOCKER. That’s the reason why you will find the special case when setting the browser arguments in summary_bot/scraper.py:

Browser arguments creation for two cases: when the bot is running locally or from Docker.

We organised our code as a Python package with setup.py.
On line #27, we install this package (in the -e editable mode because that’s how you usually do it locally) such that all functions in our source code can be imported with regular imports and become visible within the project.

On line #33 in the Dockerfile, you will notice the pytest command. By this, we run the unit tests and check the web scraping and the summarization model functionality on test data. It is a bit dirty solution, but I wanted to have a quick way of verifying that the container was built without errors related to the puppeteer.

To build an image and run the container, execute the following commands:

Build an image and Run the Docker container

Demo 🎉

After running the bot locally from the python environment or from Docker, you can open Telegram, find your new bot by nickname and insert a URL that you want to summarize.

Let’s try to summarize this very blog post about building a bot summarizer😅
It processed it for a few seconds, divided it into 3 chunks and sent me the summary 🎉

It might be a quite simple bot but I already love it since it’s nice and does exactly what was planned!

Conclusion and TODOs

We’ve covered all the main steps in building a Telegram bot that serves the AI model (in our case, the summarization model). I encourage you to try running this project and then try to replace the model with some other model from the HuggingFace repo. The Telegram Bot platform is very powerful, and you can build much more comprehensive apps with it. I like this approach because you need to concentrate only on the Bot logic, and the rest of the messaging and UI features (like various content types, buttons, etc.) are ready for you to use. Moreover, with the ChatGPT hype, wrapping an NLP chatting model into a good messenger is very natural, so please try it!

What’s next for this project? Hypothetically, it can be improved in many ways:

1. Model: play with other summarization models and model settings, maybe OpenAI models with API, but they are not free compared to HuggingFace
2. Long articles processing: currently, we split a long article into chunks and create a list of summaries. The chunking is not the best option since the summarizer loses the context from the previous chunks and considers every chunk as a separate story
3. Web Scraping: it’s a never-ending battle between web scrapers and website owners who don’t want anyone to scrape their sites. The current scraper can’t scrape the content of certain websites even with the Puppeteer (for example, the very same Medium). This part can definitely be improved.
4. More Bot features: you can add more features that the Bot serves, more logic and/or models. Potentially, you can add more commands and process various content (images, geolocation, video, etc)
5. Scalability and proper deployment: If we want to serve many users with this bot, we need to think about the appropriate deployment and scalability of the model. Should we retrieve the messages from a queue to run models in batches? Should we separate bot and the model serving at all? How to scale this system to support X users?

That’s it 😊 Thank you for reading the post!
I’ve spent more time writing it (5 days) rather than coding (2 days). This is my first blog post so I would appreciate any feedback and suggestions!