A Retrieval Augmented Generation system to query the scikit-learn documentation
Currently, the scikit-learn website provides an "exact" search engine based on the tools provided by the Sphinx Python package (i.e., https://www.sphinx-doc.org/). The current search engine is implemented in JavaScript and runs locally using an index built when generating the documentation. This solution has the advantage of being lightweight and does not require any server to handle the query. However, the complexity of the query treated is weak: since the search is "exact," it is not robust to spelling mistakes, and the search is intended for searches based on keywords.
As large language models (LLMs) are becoming more popular, we have been interested in experimenting with this technology, knowing that they could address some of the previously stated limitations. As an open-source project, we have limited resources in terms of compute and limited available datasets; therefore, we discarded the option of fine-tuning an LLM and leaned towards retrieval augmented generation (RAG) systems.
This talk presents an experimental RAG system developed to query the scikit-learn documentation. As constraints, we impose ourselves to use an open-source software stack and open-weight models to build our system. The talk is decomposed as follows:
First, we provide some background on the RAG system and the pipeline to follow to implement such a system.
Then, we go into details in the different stages of the RAG pipeline. We provide some insights regarding documentation scraping strategies that we developed by leveraging the numpydoc and sphinx-gallery parser. Then, we discuss the solution that we tested to perform lexical and semantic searches. Finally, we explain how the context found can be fed to the LLM to help generate an answer to the user query. We provide a small demo to compare queries performed on an LLM-only system and on the developed RAG system. All the code for the experiment is hosted at the following GitHub repository: https://github.com/glemaitre/sklearn-ragger-duck.
Finally, we put into perspective the gains and pains of such an RAG system when it comes to integrating it into an open-source project. Notably, we question the hosting and cost of such systems and compare it with other approaches that could tackle some of the original issues.
This session took place in track Generative AI and was classified suitable for intermediate domain / intermediate python by the speaker.
Transcript (auto)
Auto-generated from the recording utilizing Open-Source AI. Speaker labels (Speaker 1, Speaker 2) reflect diarization, not identity. Timestamps refer to the recording.
Speaker 1 [00:06]
So, hi everyone, so I'm Guillaume Lemaitre and today I will speak a bit about RAG to query basically the scikit-learn documentations. So who I am, to speak about scikit-learn, so basically I have a PhD in computer science but here this is not really relevant and otherwise I'm doing open source at Probable and I'm one of the core developers from scikit-learn and another RAG query which is eBalanced Learn. And I'm not an LLM expert, okay? So whatever I say, you have to don't trust me. But it was a fun experiment because here it's about a proof of concept, let's say. So motivation on what we were trying to achieve. So if you already went on the CYCLINE documentation, you probably saw that website. And then you have these ugly bars on the top there, which is a search bar. And if you search something, sometimes it works, so like in this case, then you end up inside another ugly page that shows you the results. And most of the time it doesn't work, okay, because you just misspelled something and then it doesn't do fuzzy matching, so it doesn't work pretty well there. Or you try to ask questions now that we are used to JTPT, and of course that doesn't work, Okay, so it's finding like terms, but it doesn't work very nicely. So the motivation was to say, okay, maybe, so when I submitted the talk was in December, so November, everything was about drag and LLM, and I was like, what is these things? Maybe it could be cool to just look at what it can give and if it could be useful basically for the cyclone project. And can we just replace this or what can we do with that? So the idea was really to go inside the directions. So first, yeah, improve the user experience. Second is to see what LLM can provide and what has the potential there, and I wanted to come in Berlin, so it was a good place to make a little talk. So I will start just by giving a bit of background on the RAG in case that some of you are not very familiar with it. And then I will go into details how we just use them for the cyclone documentations. So here the first thing that I want to mention is how you are using normally an LLM or here I actually like how we do it inside IK. So it's a zero-shot prompting in the way that you write your questions, you give this to your models, and then it's just answering to you. So it's generating one answer. We don't expect to have a chain of thoughts, let's say, where you can just ask, oh, but can you disambiguate what I'm seeing? It's not what we are seeking here. It's like one question, one answer. So what RAG brings onto these things is just that you have a preface just before that where you can say, hey, actually, the LLM is dumb. he will not be able to answer. So what we'll do is find some pieces of information to feed to him and with as well maybe his knowledge as well, can combine all of those together to give me a much better answer. So now instead of having one params, which was LLM, you have two params because you have also the search engine here, okay? So you have to work on this and on that part. So what is this information retrieval here that retrieves the context? Basically, this is not something new. It's pretty old. I mean, like, Google did search engines during a long time. So you are very close to that thing. So the idea is that I can provide the questions, and then these questions, which is in the natural language, should be brought into mathematical representations where I can do mathematics on that. And one that I have that vectors, I could have like a database and check what are mathematically what are like vectors that look like this and retrieve the vectors and then go back to the natural language, okay, that give me like basically the text with that. So you need to build the database and you have different way of doing the representation here. So I will go into details what here we use as a proof of concept because you have different methods and these fields have been quite active recently basically with the language models. So now that we have kind of a framework on a rag and that part, let's build a rag, okay? So if I go around, so I did the slides two days ago, I just like how to build a rag and then you get a bunch of pages and it seems super easy, okay? So one does it in 50 lines of code, another one in four lines of code. that you can just put it in production and it works. And Mistral, they do basically models, so they tell you all the different ways that you can do it with line chain, Lamia index, or XTAC. So it seems like this is something straightforward, and you should have any caveats, okay? Because I mean, like, all the different people can tell you how to do it. So what is also interesting is that if you check, like, the different blogs, the first one is Scaleway, so they sell you some GPUs, AnyScale does the same, Mistral doesn't sell you but they will give you some LLMs, and Yamaindex and Clarify in that case give you some basically right out of the box. So all these people have an interest of course. So here I will give another bit, another perspective on the PROMs, linked to my committee work let's say. So I will put some rules compared to those people. One is I'm kind of poor. I don't have GPU. I have a Mac M2 that can do stuff. So I want to work on my machines to develop those. I'm stingy. I don't want to pay any of the providers. I don't want to bother with API calls because I could still pay OpenAI and just make API calls, but I don't want to bother with that because I'm still stingy. I want to write pure Python code, I'm a cyclion person, so maybe if I can do my own wrapper, I can do that with the cyclion API. And then, so this is not a given, but if you know bitlangchain or those libraries, this is wrapper of wrapper of wrapper of wrapper, so I can avoid that because I want to understand what is under there. I don't say don't use them, I just say I want it to run what is under that. Okay. So that's the thing. So it is called Raggeduck. So it's as good as a rubber duck debugging because it can maybe tell you wrong stuff but maybe it will tell you the truth. And it's Powerpack Pixie. There's a talk tomorrow on Pixie. Go and check. This is super cool. So here there is some documentation. If you want to reproduce locally this rag, actually you will be able to do that by following the documentations of this project. It's done to be like reproducibles, okay? So what choice did I do for that specific proof of concept? Of course, any of the blocks here can be replaced with any of the technology. That was not my interest to pick up the best things. The only thing that I wanted is that it works and fast. I can pip install it, okay, more or less that. So I will come back on the retriever, but basically I have a sentence BERT that comes from SBIRT, FACE that is a library from Meta to make search, nearest neighbor search, approximate nearest neighbor search. I implemented the bm25 because there is no bm25 in Python apparently. And then for the LLM part I'm using Mistral because I'm French. And I'm using YAML CPP basically to have the interface with it. And then I found someone that writes web things to be able to question the things basically online that is called LemaWeb. And it was more than enough for what I needed, so I grabbed a piece of that. I mean, that's a nice part about the open source. And I just throw it in there, and then I have a full thing that allow me to interrogate basically the LLM with the context. So now let's go into the details because the details are like that's the devil in the details. So if you look at all the documentation that I showed before, they will say, yeah, you take a text, it's easy, I mean, you document that you want to retrieve, I just chunks of documentations. So if you define it, it's like I did that page, I take the first 300 tokens and then then I take the next one, and I can overlap between them. That builds basically all your documentation that you want to make search inside. And at the moment of the retrieval, you will just actually pick up one of those chunks and provide it to your LLM. So you say, OK, good. I will start with that. But we have different type of documentation in scikit-learn. So maybe I can try a different strategy, and as well learn what is wrong with this. because actually, if you try that, that doesn't work. So the first type of documentation that we have is the API documentations. So it's a reference one. So we have classes. We have parameters. We have defaults. We have the description of the parameters. So you can imagine the kind of, basically, question that people could come with when they try to search inside those documentation page. So these things like, what are the parameters of logistic relations? Or what is the value of these parameters? Or give me a piece of code that give me a use example, because we have a section with that. So it really is the type of question that people will come. Or let's say this type of documentation are a good candidate to answer those questions. I don't say that that's the only one, but this is one of them. And I will try to show you that if I use the previous strategy, which is a naive one that's like 99% of the text, like documentation tell you to do, is actually failing, so I will try to answer the second one and give like a content examples of how Wrike failed there. So if you try to make the search, basically in Scikit-learn we have different places where you have the strategy one, and they will come with, I take like the three most important chunk, and if I chunk in this way, I will get like a first chunk that is coming from dummy classifier, which answer perfectly the questions. And I have all the details on what are the possible values. If I take the second chunk, that will be as well provided to the LLM, because this is the most important one. Then here, this is coming from the dummy regressor. But I have nothing in the text here that allow me to know if this is the case or not. So I really bring some ambiguity to my LLM, because I provide him stuff which there is no information to disambiguate basically what is the sources, for instance. And even worse, so classifier regressor, I mean, you could expect, but the simple inputter in section as well have a strategy. So this is, again, something that doesn't work. So now if I take my models, I ask the questions, and I provide him the question plus those context chunk of data that I have before, He will answer you these things, which is like good hallucinations. He will tell you that dummy classifier doesn't take any strategy parameters, OK, while you provide him the right things. However, he will tell you that you speak about this with a typo, which I have no idea where it comes from. But I mean, he just like create a typo, and then speak about like missing values, OK? So he's just like certain of what he said. So, and you just felt on the specific questions because you bring ambiguity inside the different checks and you just try to get out with it and just give you a bullshit answer. So this is really problematic because it's not what you want to provide to your user. So how you could do to try to disambiguate that? So in some ways, the documentation in Scikit-Learn is organized in, we use NumPyDoc. So we have a way that we can scrub the documentation in the structure manners, and we can recreate, basically, chunk of data, our own text, that use the previous text, but that adds the information that we know that could be ambiguous. So in that case, I could always, for instance, say, what's the class we are speaking of, and then generate a text with all the parameters, and say, oh, and the default value are those one, and generate different chunks. So for instance, this is exactly what I did for the different classes, the different functions. So you write a smarter way to parse the data, OK? And if you do that, then you get the following answer, which is a much better one, and so I will not read it. But basically, this is the correct one, which corresponds to the first chunk. Because in that case, the context saw that dummy classifier was very close to this strategy compared to the other one. So the LLM in that case managed to use these ambiguations. If you ask me why, I have no idea. But it kind of works. So that was for the API documentations. Then we have two types of over-documentation. One is a user guide, which is, let's say, it's explaining some mathematical details and stuff like that. So it's going further than just programming. In this case, I used the naive chunking, but we could do better. We could detect basically sections and be sure that we don't cut chunks outside of the section to be sure that basically what we have here is consistent and that you don't cut in the middle. I didn't implement that, but it should be done in some way. That would be much, much better. And we have a gallery of examples. So usually these are examples. There is two kinds. You take an example like on the left, so you have just a bit of text and then codes. Over here you have tutorial-like examples, so different sections. So for that one I use Sphinx Gallery because it's parsing actually the data and I can just be sure that a section itself is not cut it, so I can regroup basically section together and then generate the chunk from that. While in the first case we don't care because you have only two block of data. So this is what I just described here. So on the tutorial one, we just detect sections, merge the sections together, and then chunk from that. While in the second case, we just chunk every individual section, like the two sections. OK, so this is how basically the data, let's say the data, the documentation was scrapped. Now what we need, to go in a bit more details, is on the retriever. So the thing that will allow me to find the right chunk. So as I mentioned, it's something which is very old. And you had word-based retriever a long time ago. And it's what we call lexical search. And it's something that's so when you have your document that comes from on the top, you can basically count how many times a word appears. So modulates a bit of a refinement. But basically, you get a sparse matrix with the counting. And then you will get your query where you check how many times the words are appearing as well as the document. And then you can compute a score. The score will tell you how likely two documents look the same. And then from there, you can just extract the document. The problem of this approach is that it doesn't understand the meaning of a sentence because you are cutting out words. So if there's a word, it doesn't understand if it's negated or not, okay? So it's just like, this word appear. So it doesn't understand that, but this is very efficient, because this is just an inverse index to retrieve basically the document, so that's working very well for that. On the other hand, with language models, then people start to use transformer and BERT models to actually take documents and project them into a dense embedding, and in that case, because of the transformer and the attentions of the transformer, you encode context. So you encode in which context a word appears. And this is kind of very interesting, because that's solving one of the problems of the one. But then you have a problem with the retriever, because in that case, you need to use something like approximate nearest neighbors to find what is the closest vectors there, and then to go back to the document that was linked to your vectors. So it's a bit more involved, I mean it's a bit, for instance, it will be ten times slower than 100 times slower than like the BM25. But it's understanding context. Okay, so here this is just like to show some code because we are in Python conference. So somehow I wrote some scikit-learn compatible estimators that allow to scrap and then you you can use, for instance, a pipeline and just put the thing that will just extract the text and then train your models. You call fit on the directory where you have the documentations. And then you just have to pickle, not the extractor, because you don't need it, but just a retriever, because we'll need it in the next stage when we want to make the write. And the same exact thing for the sentence transformer. So it is just like, I have a cyclone interface there. I call fit, and then I dump the pickle file. Because we have two type of retriever, we don't know exactly how to order them. So there is an extra step, which is a reranker, which is, again, based on BERT. And the idea is that you give the query a document, and it will give you a score saying how likely this is. And this model was trained on BigQuery. So somehow, I don't know if this is the best here. I mean, I will come back to this in the evaluations. But in some way, this is a good start to say, OK, I find all the documents. I want to re-rank them. And then I want to take the fifth best document to answer my questions. And the final piece that is missing is I need, so I can now provide the context in my LLM, but I need to prompt it. And there is nothing fancy there. So I just ask the typical thing that you say, that you are a very nice expert in whatever, and then answer to these things using this context. And when I extracted the context, I as well extracted the HTML page sources. I said, okay, use those two information, and then go on. So if you put this all together, that looks like that piece of code. So this is, let's say, the web app where first you load all the retriever that I dumped during the training. Then you create the re-ranker which is take all those retrievers. I create my LLM afterwards with YAML CPP and then I provide the LLM plus the retriever to the prompting that basically will be in charge of just like putting the informations. So, demo time, and because the internet is not working very well, I just make the prompt before to come. Okay, so basically here I ask exactly the previous questions, the answer is correct. The only thing that could be interesting in that case is on that page. So basically, this is on the back end, so the first thing that we do is that here we are asking the different retrievers to find the chunks. So this is the time that happens, and for instance, on this machine it takes three seconds, so this is pretty slow. If it's on the GPU, this is more closer to a second, let's say, but it takes, let's say, In total, something like five seconds to find the chunks, then the prompt itself looks like that. So here you have the answers, you give the query, and then so you're asking, like, you are an expert, then answer that query, and then you provide the context, and the context is the source of the documentation and then the different chunks. you provide this to the LLM, and it answers you back, basically, that things, okay? And in that case, is where it's working quite well. If I disable completely the LLM on that things, it just say, like, very, like, it never saw documentation of cyclone before, so it will tell you two parameters completely random that you found somewhere, okay? So the RAG is still working this way, I mean, like, whatever you give him as information, LLM will use it. So time for conclusions, I think. Yes. So part of the evaluations, one problem here is that there is no evaluations. Why? It's because we don't have data as a community-driven project, because we don't collect data on the website, so we don't collect what people type on the search bar, so I cannot use that to have a feedback or to fine-tune my models or to know what works best. And that's fine. We don't want to have that. And for GDPR competency, we didn't want it as well to go with that. But it means that it is impossible to evaluate our right system and it is impossible to fine tune it. And some other, because that's my last slide, is where I think that we need to be a bit critical of the technology. So first as a community project for scikit-learn, we have kind of a million users per month on the documentations. If everybody starts to do a couple of queries, you see the numbers. It's costing zero, it's costing kind of like a couple of cents per hundred queries. We compute it, we need 10,000 euros for this toy, okay, per month. So we cannot pay for that. So if you're inside a company that's making people inside the company very effective at doing their work, this is completely fine. Here, this is completely another setting where it's probably not valuable for that type of work. However, what we can do is a much better search engine, which is fuzzy, let's say, fuzzy terms. And Algolia, for instance, is proposing something like this to open source programs. So that would work. And that's it. Thank you.
Speaker 2 [24:25]
Time for questions. Do we have questions? No questions? No questions. Okay. I have a question. So some of the tutorials contain images, and how do you process questions where you need So, yeah.
Speaker 1 [24:43]
So here we don't process the image as simple as that Is a text base, so then then yes if we would like to push until until that stage then you have no work that You are I mean you can understand like the image maybe one way would be from the image Can you train a model that can generate the code that? Generated the image and then like use that as a context as text as well Or I don't know there's like multiple strategies that we could think of But here, yeah, I didn't go there.
Speaker 2 [25:20]
Thanks for the nice talk. So it was just better more to understand if I got correctly your chunking strategies. So basically, what you recommend is for every different kind of document, you have a different chunking strategy.
Speaker 1 [25:36]
Yeah, you need to, so first is like what I saw is you check how it fails and from what it fails you try to understand why and usually you will come up with a much better chunking strategy and the default strategy is not going to work all the time, for most of the time. So yes, you need to spend, the main message is if you put garbage in you will get garbage out. So really pay attention to what you put in and actually then come with a better strategy for chunking documents. And it's case to case. I don't have a solution here to tell you, oh, for you, you need to use that. It will be case to case. For instance, for the user guide, the naive chunking is already a good default. For the API, this is not at all working. So this is the kind of thing that is interesting to know and that people don't tell you.
Speaker 2 [26:29]
And then afterwards if you would have let's say many different type of chunking then you assimilate them all by the Reranking right and that is a way to just
Speaker 1 [26:38]
Yes, the re-ranking will just reorder depending on what is the best. I mean, you already find the right things.
Speaker 2 [26:50]
So what I ask myself is you're doing your retrieval, you're announcing it retrieval with BERT and then putting it into the LLM. So wouldn't it be good to just check the retrieval first if it's gotten better? Because that would be much cheaper, I think.
Speaker 1 [27:13]
So, for me, the LLM here is, I mean, if I'm very, like, critical and mean, the LLM is useless, okay? What we need is what we are actually searching as a product at Cyclone is a very good search engine. And these things, yes, we could evaluate it, but to evaluate it, we need to collect user data. We need to collect what people are searching, and if they are clicking on the first link, is that we propose the right things. And I would say, if you want to do a good RAG, yeah, evaluate that part, because it is critical. It's what gets in. And I agree with you, and that might be cheaper. And you have a better feedback, like, as your feedback in that area, than like, evaluating the output of the LLM.
Speaker 2 [28:03]
So you mentioned performance of the context retrieval So what was your experience with mistral on CPU and the number of chunks that was reasonable to run on?
Speaker 1 [28:15]
So, first, like the LLMs are, for instance, the last ones are limited with the context, so the context cannot be more than 8,192, and the first version was 4,000. So you can tweak those numbers on CPU. I would say this is a dead end, it's super slow. On the MPS, however, this is reasonable to, I develop everything on the MPS, so I mean on the mic M2, so it's using like the chip, so it's quite good for prototyping, and on the GPU this is like, that's a normal speed like what you see in ChatGPT. Do you have to wait a couple of seconds or tens of seconds or minutes? Let's say it will take 30 seconds to answer one query. So it's OK as a latency to develop a POC. It's not OK to have it in prod. In production, you will need a GPU.
Speaker 2 [29:23]
Okay, thanks so much everyone for your time. We have now come to the end of this talk.