Transcription
[Music]
We are from the AI development and implementation team at cloud.ru. My name is Pyotr Korolev, and I am a Data Science developer. With me is my colleague, leading Data Science developer Mikhail Dromi.
Before we discuss the recipe for creating a RAG pipeline, we need to clarify some terminology. RAG, or Retrieval-Augmented Generation, is an architecture or approach to information retrieval that combines two key technologies: information retrieval and text generation.
On the slide, you can see the standard architecture of RAG. Let's take a look at what it entails.
First, the process begins with a user query. Next, using vector similarity, we select the top K relevant paragraphs from our document database. These relevant paragraphs, along with the query, are then fed into a generator, which is often a language model (LM) that generates a response based on the retrieved snippets.
Why not just train an LM on our data and make it answer questions from our documentation? The reason is that fine-tuning LMs can be very expensive and time-consuming. This issue often arises when your data source is frequently updated, such as documentation for your product. The RAG pipeline circumvents these limitations by providing current and relevant information for generating responses to user queries.
Below is a more advanced RAG architecture that is commonly used in production. Here, two new components are added: first, a neural network-based retrieval system, and second, a ranking system that ranks this mixed output.
Now, the clicker seems to be refusing to work.
So, why was this important for us? Well, we have three cloud platforms and dozens of cloud services, and consequently, the documentation for them is quite extensive. Searching through this documentation can take some time. The implementation of the RAG pipeline helps reduce the time spent searching for useful information and increases user loyalty. Additionally, it reduces the workload not only for our users but also for our technical support team, which responds to user queries about the documentation.
Thus, our task was to integrate the RAG approach into the search functionality of our technical documentation. The first step in solving this task is data collection and annotation. We will collect and annotate data for two main components: the retriever, which is responsible for retrieving pieces from our document database, and the ranker, often referred to as the scorer, which processes the search output.
Data can be collected and annotated in two ways: synthetic data and assessor-based annotation. In an ideal world, it is best to train everything on assessor data, as it yields the best results. However, such data is often scarce, and that was the case for us.
We annotated the data and then passed it on to our experts for further annotation, having prepared instructions and developed an annotation methodology. As a result, we obtained a dataset consisting of queries, texts, and relevance scores ranging from zero to three, where zero indicates irrelevant text to the query, and three indicates super relevant.
After we submitted the data for annotation, we began selecting models. Our first criterion was to look at various benchmarks, primarily focusing on the Ditch benchmark, which rates Russian-language encoders.
Next, good tokenization for the Russian language is important. On the slide, two histograms are presented. For example, the right histogram shows the distribution of the number of words and tokens for the best model, specifically the MiLM, which is recommended in the most popular encoder rankings. It has the best score, but as we can see, its distribution is poor, and its context length of 512 tokens accommodates only about 100 words.
On the left, we have the XM-Roberta Base model, which was trained not only on English data but also on Russian data. Its context length accommodates around 300 words, which is much better.
It is important to remember the context; this distribution is crucial because we need to fit as much textual information from the documentation into this context window for proper ranking to work.
Now, let's talk about how we trained the models. First, we focused on training the cross-encoder, which was later used to train the encoder. The encoder is a standard encoder with a classification head that takes two text sequences as input, separated by a separator token. The first text sequence is the query, and the second is a paragraph from the documentation. The output is a relevance score ranging from zero to one.
Initially, we trained the encoder on the RAG dataset, which consists of a machine-translated version of the MS MARC dataset. While it is not suitable for final model tuning, it serves well as training data.
Next, we conducted supervised fine-tuning (SFT) on the annotated data, augmenting it with various errors.
Now, let's discuss the training of the second model, the BERT model. The main difference between the two models is that the cross-encoder is quite robust across different domains. The Russian encoder sees both the user query and the text to be ranked simultaneously. In contrast, the BERT model is somewhat "blind"; it only sees the user query and computes the relevance based on it, only seeing the relevant fragment.
It is also worth mentioning that the BERT model is not very robust across different domains. If you take a pre-trained model in your domain, it may not perform well. To conduct domain adaptation, we used the Generative Pseudo-Labeling (GPL) technique. Essentially, this is synthetic training, where we take snippets from the document and generate five to seven user queries using an LM.
We then use these queries to find search results in a current search system, such as OpenSearch or Elasticsearch. The generated queries can be sent there to obtain hard negatives, which are sufficiently relevant paragraphs that will be used in our contrastive learning.
We obtain a triplet: the generated query, the positive paragraph (the one we retrieved from the search results), and the negative paragraph. Next, we perform pseudo-labeling using the cross-encoder model. Since it is not heavily biased towards any domain, we can use its scores as a form of conditional ground truth labeling.
We take pairs of queries and paragraphs and evaluate them with the cross-encoder, resulting in a quartet consisting of the generated query, the positive paragraph, the negative paragraph, and the scores for the query-paragraph pairs.
Next, we conduct standard training. We input a positive example and a negative example into the encoder, obtain embeddings, and then use margin M loss. This loss function optimizes not a specific target but the relationship between similarities.
As we can see below, we have similarity scores for the user query and the positive example, as well as the similarity score for the user query and the negative example. We take the difference between these scores and optimize the margin. This allows us to approach training flexibly and be tolerant of the presence of, say, good examples in our current search system's output.
Hard negatives can also be positive, but the margin between the positive and negative will be quite small according to the cross-encoder's evaluations.
While GPL is beneficial, it does not solve all problems. Since we are in a fairly specialized domain—user documentation—there are specific terms that users frequently search for, and good recall is crucial. For example, if we take a narrow topic like nuclear physics, and a user wants to search for quarks, the model must understand that quarks and elementary particles are essentially the same. However, the GPL technique sometimes fails to achieve this.
Therefore, in full-text search, we use a synonym dictionary that we obtained from our technical writers for these specialized terms. For instance, the term "virtual machine" can be referred to in various ways: virtualka, VM, virtual machine, server, etc.
After adding this synonym dictionary, we also implement filters to improve full-text search. These include basic filters like converting to lowercase and removing unnecessary parts of speech, such as prepositions and conjunctions.
The Russian morphological analyzer (Rus SBO) helps us be somewhat resilient to grammatical cases by trimming the ends of words. We also use a shingle filter that allows us to search for complex terms like "virtual machine" or "distributed attack." We search not only by unigrams but also by phrases.
Here, you can see the results of our training. The main goal of our presentation was to highlight specific problems in a narrow domain and propose concrete solutions.
We measure the quality of search results using metrics because we input exactly three texts into the model to generate a response to the user query. It is also possible to input more or fewer texts, but we found that when we input more, the model often gets lost because the texts may be unstructured relative to each other or come from different documents.
Thus, three paragraphs yielded good results, and we decided to stick with that. For example, looking at the second line, we used a good output obtained from the Muling E5 LGE model, which is a solid encoder with good tokenization from OpenSearch. However, the result was worse than just using BM25 because the MiLM essentially cut off 45% of the text.
The techniques we discussed gradually increased the score, with the most significant improvement coming after training on the assessor-annotated data from the Russian encoder. This approach combines full-text retrieval and semantic retrieval, yielding good recall.
However, there were questions regarding ranking; we achieved better ranking quality.
Now, let's discuss the most crucial part: the production architecture. It is great to train models and so on, but it is important for them to work somewhere.
This is what our architecture looks like. We have discussed certain parts of the architecture but have not touched on important surrounding processes, such as the ETL process for document processing.
The main challenge is that those familiar with RAG architecture know that it is quite difficult to devise a good text segmentation strategy. Essentially, we always have a choice regarding how large our text will be and how we will segment it.
If we segment it, we might cut a logical part in half. We can segment by counting characters, sentences, or tokens—there are many strategies to consider.
The first thing to note is to check the results of open-source libraries. There is a popular library called Long Chain that is widely used in various how-to guides. We encountered issues early in the library's development when we tried to segment our text using regular expressions, and the local splitter did not segment according to those expressions but simply inserted them into the text.
So, one should not blindly trust everything. Additionally, if you plan to use this in your production code, it is important to understand that this library updates rapidly.
In general, one of the best strategies is to segment text by logical blocks. If you have HTML markup, for example, you can segment it by tags or paragraphs because people typically use these tags to logically separate one thought from another.
In the end, we use a strategy where we recursively segment by various H tags, from H1 to H6. If, after segmenting by the H6 tag, the text still does not fit into the context window, we further segment by paragraphs.
Among all the strategies I mentioned, it is best to count the length of our chunks in terms of tokens rather than words or characters. This is because we have two models, and it is clear that one can tokenize more words into a 512-character context window than the other. We take the smaller result to ensure everything fits, so that a domain-specific term does not get cut in half because we counted characters instead of tokens.
To manage all these processes, we use Palai, which is very similar to Airflow. Its main feature is that it is very developer-friendly. Essentially, you do not need extensive knowledge of how it works.
In our case, we figured out how it all works in one day, thanks to its rich documentation and friendly code.
To update our documentation or any knowledge base, we can use two approaches: either schedule tasks or have a webhook. Initially, we considered using webhooks to trigger updates based on events, but it turned out we were not ready to upgrade, and webhook functionality is only available in the Perfect Cloud version.
In the end, we decided to run scheduled tasks, executing them about once a day at 5 AM to avoid disrupting user experience.
The process looks like this: we have a GitLab pipeline where technical writers store all rendered documents. We take these documents from GitLab, version them according to the described strategies, and then save the resulting texts in the Milvus vector database.
Before saving anything in Milvus, we use our encoder to obtain text embeddings.
As for deployment, everything is set up in a Kubernetes cluster. It is important to understand the deployment paths for the described components. We deployed OpenSearch using the OpenSearch operator, which is a native approach for Kubernetes.
Based on our experience, if you need to deploy OpenSearch in Kubernetes, you will encounter fewer administrative issues, such as password changes, compared to deploying it as a vector database.
For the vector database, we used Milvus, which is quite mature and has many stars on GitHub. It is also cloud-native and supports horizontal scaling of components. It has a microservices architecture, allowing you to tune each component elastically under load.
Milvus offers better performance compared to other vector databases, such as Quadrant and others. It also has SDKs for all popular languages, including Python, Java, and Go.
To serve models, we use the NVIDIA Triton Inference Server, which achieves very high requests per second (RPS) due to its architecture.
In principle, if you want to go further, you can also quantize models to 8-bit or 16-bit to enhance their performance.
The user experience looks like this: the user enters a search query in the documentation portal. This query goes to our search component, which implements the basic search logic.
First, we query OpenSearch to obtain full-text results. Then, we retrieve the vector of the user query and use it to obtain semantic results. We simply mix these two outputs—essentially adding one to the other—and feed the combined input into the ranking encoder.
We then obtain the final output and send the top three paragraphs to generate a response to the user query.
Did we succeed? Yes, we did! The results are presented on the slide. This is how it currently looks in production.
Thank you all for your attention!
Thank you, everyone. We have a couple of minutes for questions. Are there any questions? I see a hand raised.
Yes, thank you for the presentation. Could you tell us why you chose to tackle the problem using the RAG approach? There is also the RAG in QA, which seems more suitable for instructions.
That may be true, but some of our questions are quite dispersed throughout the documentation. There isn't a single place where the text is very concise. Before we used RAG, we tried various summarization models, and they performed poorly in our case, which is why we opted for RAG.
Thank you. Hello, thank you for the presentation. I have a question: you showed a user evaluation on the slide. Do you have a test dataset on which you evaluate, and what results does it yield?
Yes, there was a slide with a green table, which shows the evaluation on our synthetically generated queries. We then submitted this synthetic generation for annotation by assessors to make it easier for them.
So, the assessors checked the results themselves? Yes, they did not annotate from scratch, which saved time and increased productivity.
Thank you. Yes, thank you very much for the presentation. Can I ask how you combine BM25? I understand that it is the formula by which the relevance of the query and document is modified.
Yes, that's correct. We have two outputs: the full-text output obtained from our vector database. Initially, we simply concatenate them—essentially, we take the two lists and combine them.
However, we also have this module, which we referred to in the presentation as a model that takes this dual output and processes it, assigning final scores. Essentially, it further ranks the results from both outputs.
Thank you.
[Music]