Scaling Embeddings in Real-Time

SBTB 2023: Zander Matheson, Scaling Embeddings in Real-Time.

Recording: SBTB 2023: Zander Matheson, Scaling Embeddings in Real-Time.

[Music] or something like that we're going to talk a little bit about shoving things into Vector databases here but we're going to talk about how you can do it in real time um and in Python uh and we're going to use bite wax which I'll talk a little bit about in a second and I am more than Loosely affiliated with bite wax and that's why I'm wearing this uh bque shirt here so we all love the new Power of large language models um probably all been playing with them or building things with them um and we've probably all noticed that they do have a little problem um for example I want to find out uh with the um new Barbie movie gross and who were the main actors but unfortunately that was after the cut off um so the problem is is that the language models themselves are stale stale isn't bread nobody likes stale bread if you serve stale bread to your customers they will leave if you serve stale data to your customers they will also leave so in our instance with the Barbie um the Barbie movie question we would probably leave chat TBT well maybe not because of future changes but we would leave chat GPT or we'd leave the large language model application and we would go back to Google to find out what happened now I know what you might be all thinking and that I made this before the open AI Dev talk and rag is obliterated etc etc um it's I don't think it is but I just wanted to address this uh in the instances where you were using chat gbd to do that it would look up in beinging and provide you some context and and answer that and we'll talk about that in a second but in most applications specific to the large language model itself you're still going to come upon this rag problem so we'll ignore the um open AI Dev day for for right now so a little bit about me before we start uh continuing on here uh so I work on bite wax um if you want to head over to bite waxite wax and give me a star that would be amazing thank you um I'm also sitting right in front of this room for the rest of the conference so if you want to talk about anything or if you want me to make you a PO over coffee I'm making por over coffees uh because conferences always have really bad coffee no offense to the to the conference organizers whom Oli who is part of my team was a part of it's just a fact uh when I'm not working on bite wax I'm also a uh dog and now a human Dad if you are starting a new company don't have a kid it's really hard um outside of work I like to go skiing and generally get Outdoors this is me um mostly trying to look cool but also trying to go skiing and when I'm not in the bite wax slack um where I spend most of my waking hours um I'm in Santa Cruz and if you're ever in Santa Cruz you want to grab a coffee hit me up on LinkedIn all right so back to our problem about stale bread I mean stale data um how do we solve this today what this is what we're going to talk about today so we're going to talk about solving that problem we're going to talk about how we can do it with retrieval augmented generation we'll talk about what a little thing called an embedding is and then we're going to do a technical implement of how we would make a pipeline to have uh to serve real time uh data in a retrieval augmented generation system uh powered system so to solve our stale problem we have three ways that we can kind of do this one is retrieval augmented generation where we're essentially finding the most upto-date or relevant information about the question being asked and we're shoving that into the model prompt and then the model should generate what we want um you could also fine-tune a model I mean all of you probably know all this you could fine-tune a model with upto-date or relevant information so that it is available for retrieval in some way shape or form you could also retrain the entire model and so when you see the cut off that I showed in that photo it was like a new uh model that had been trained in with the cuto off I think of April or something like that in 2023 which prior I think it was like 2022 or 2021 anyway um if you know about the models they train or they retrain uh it's quite expensive and takes like way more um compete resources than I have at my fingertips so retrieval augmented generation is what we're going to talk about today um and how you too can get rich building rag powered AI applications no not not actually but I just I don't know I thought I was being a little bit funny by making this headline so we're going to talk about rag so what is rag how does it work as I kind of touched on you're going to get a query hey I want to know something about something or explain to me something and you're going to go out and you're going to retrieve what you um hope is the most relevant information then you're can going to combine that query and sort of instruct the model to say hey this is the context you have please answer this then your large language model would provide or generate the response and you'd send that back to your user so in order for us to go any further here we need to talk a little bit about embeddings so we want to be able to store some information in a retrievable place uh whether that be a vector DB or not uh one good way to and we want to be able to compare a question that's being asked to this information that's stored in a reasonable way and retrieve that so that we can then shove it in front of the model uh one really nice way to do this is uh the embeddings so embeddings are a way to essentially take things and turn them into numbers so in our case we're taking sentences or you take words or sentences or paragraphs and turn them into a um an array of uh numbers between zero and one so here we can compare um a cottonade maritime shirt which is apparently really close to a striped blue shirt made from cotton I don't know what they would call my shirt but I don't think it's a maritime shirt but the uh embeddings allow us to compare how close those are in the vector space so we can more uh easily retrieve the data okay so adding to our little diagram we saw before um in order for us to retrieve uh we could use Bing search or Google search or some other search mechanism or we could uh with the user's request we could put a document in front of it and that's sort of what the new capabilities released in the open AI devday That's the approach you would use there but what we want to have here is a separate like asynchronous pipelin so we're not like going to force the user to provide us with the context before we answer them we want to be able to have all that context in this Vector database so we're going to write an ingestion pipeline where we're going to be retrieving data from a realtime data source or near realtime data source um one that is topical um and we're going to put that in a vector database and then in theory we' be able to make the rest of our AI application where we can take the user's query um get create embedding from that and then use the vector database to compare how similar the information is that we want to retrieve from Vector database and then we'll put it inside in the context um put the context inside the prompt that goes to a large language model um for this uh example that I'm now going to go to the more exciting part where we get a look at code uh we're going to use a um uh data source where you probably all um either look at it on a regular daily daily basis or you maybe contribute to it who knows so maybe we'll be looking at some of something that you've recently put on it and it's Hacker News so Hacker News has a really fun API if you've never used it um where you can actually get uh well you can do a bunch of things but one of the things you can do is you can request the max Item ID so we're going to kind of turn this API into a stream of um new storage and new comments and then we're going to analyze those so this is where the like real time part comes in it's not necessarily that all of these are new real time data sources but they are going to be trending in some sense we're going to use that to our advantage and so we're going to take all of that information we're going to put that in the vector database and then we would be able to query you know say what is the cool new trending database and it would probably tell us something about a vector database or something um so uh two to do this I'm going to now introduce bwax because I think this is probably a good time to do that so biox is a stateful stream processor um and it's a d data parallel by Nature so what if you don't know what that means uh essentially it's a way to process data where you decided to do things where the data um is split and then is worked on um independently instead of all the tasks are split and then you send the data there but more importantly for us it's a python has a python native API so we can use the whole Suite of uh tools that's provided in the python ecosystem um it runs on top of a rust engine which is uh called timely data flow and we've kind of built on top of that a a bunch of things you need to run this in production with a uh stateful stream processor you can do fun things like window agre aggregations joins Etc um we've focused on trying to make something that was really extensible and also uh um developer friendly uh so it's pretty easy to get started just pip install bite wax and and then you're ready to go um we've created a uh input API so you can basically connect to almost any data source and you can still get some of the things you might eventually want like uh being able to partition across or being able to split the input across multiple partitions being able to um recover from offsets and those sort of things Okay so at the very highest level uh when you write a uh bite wax data flow is what we call it um so when you use the API to describe an input some Transformations that are happening in operators and an output you've created a dag so that's what we're going to do so we're going to make a directed graph where we're going to get some input from Hacker News and then we're going to do some some initial processing uh then we'll use some of the Python ecosystem to process further process the document create the embeddings and then we'll put them into uh redus in this case uh we're going to use that like vector capability of redus um so the part of this that is the bite wax part because I I've had this question before is um it's the API to describe these so there are a bunch of operators The Operators basically um control the flow and then you use Python code inside the operators to do the Transformations so you could think of these nodes in the graph they're all operators um where you could do things like split split or join or or Windows Etc or just a map and then you pass python code to do the actual transformation so I mentioned about our input API so we'll just take a quick well a quickish look better look at the time um I want The Hacker News input uh class that we're going to build that's going to be our input connector um that's going to pull hacker new in real time time we're going to take advantage of that Max item ID so essentially oh sorry I need to go to code here uh okay is this easy to see what happens if I click that there we go okay so I know it's sometimes you're like oh what am I looking at what am I looking at so um when you're writing uh connectors for uh some framework to process data you need to be able to describe how the data is split up and then how you want to recurringly ingest the data and so with bite wax you have two parts here so one is your um partitioned input uh class and that's where you describe like hey worker one you're going to look at these parts hey worker two you're going to look at these parts uh and so that's called initially and then that Returns the next part which is the actual data source part and that's where you instruct by Wax hey every X number of seconds or milliseconds um I want you to like pull this thing and receive this data um and then when you receive that data shove it Downstream and then I'm going to do things to it uh so for in this instance we're um we've basically created a uh our our class or like h& input class will receive like a um where is it uh an interval um line two and then an initial ID and we'll ignore the initial ID but the interval is essentially the polling mechanism so if you ever use Kafka or other systems where you're consuming from a a a streaming system you're basically like repeatedly pulling and requesting data uh the new data that's arrived so in this case we're pulling every 15 seconds and then we're going to get a new Max ID and then we're going to say hey I've received this new ID the last one was this ID I want you to send Downstream all these IDs and now we now have a stream of IDs we've received Downstream there's a lot of uh streams there okay so we've written this little Hacker News input or we've described how the input looks for this Hacker News um input connector and now we're going to do some initial processing so what we've received at this point is a stream of IDs and you're like that's cool we can't do anything with the stream of IDs it's not going to make our rag application so we need to do something interesting here to enrich and then we can then do some pre-processing on those IDs uh I did have them as I think it's this one okay so before I go any further um I mentioned that things were you had this idea of an operator and you describe this like flow of data through a data flow um and so that's what this looks like here and this is that dag I mentioned so we have a number of steps where essentially describing what we want to happen so we're going to have like an input then we're going to redistribute it across workers um which I won't get into right here but basically more workers than you have partitions you can now split that across as many as you need so in our if we have Network constraints or whatever we can have many requests going on um we're then gonna oh oops hopefully there's a cach there okay uh we're then going to use a map operator and this is that like pre procing I was cutting out here so we're going to do some enrichment we're going to we have if you remember we have the stream of IDs coming through now we're going to take the ID and we're going to get the actual metadata Associated what's the URL who wrote it Etc and then we're going to take that and we're going to filter out to just stories because if you've used the API before you get polls you get comments you get stories we only want the stories right now uh and then from that story metadata we have the URL and we're going to receive or then we're going to request for the URL so that's what's Happening Here um and I I won't get into all the details here but you can just imagine that hey we've successfully received the stream of IDs and now we know which ones our stories and we're going to go get the URL and now we have this blob of HTML and you're like that's sweet I have a 60,000 token like HTML blah what the hell am I going to do with that so we have to write another little capability into this which we're going to talk about next that way can process that HTML so if you've ever worked with raw HTML from like a request of a web page it can be just an absolute pain to deal with you'll have like a bunch of jav JavaScript stuff like stuffed in there um weird HTML tags all these CSS classes they made up and it makes absolutely no sense and so we're going to use unstructured because one of the things that's cool in my opinion and I'm not biased at all about BX is you can use that like entire Suite of python uh of the Python ecosystem and there's literally a library for everything in Python on especially when you talk about data so unstructured .io has a package called unstructured um and I think unstructured and they have some really cool functionality you can do do here where you can essentially ask them to like split up the HTML and then they'll do some work and they'll like tag different things uh into different like texts so you can have narrative text you can have other descriptors on it and so we can return that narrative text only and we'll basically just get the the stuff we want the next thing we have to do is chunk that up so we can use them you can you can use UNR structure to also like chunk that uh data up and the reason we need to chunk it is because the data could the the like text you end up with after you par this HTML if that's bigger um bigger than the like uh the window or sorry the size of the embedding so you can't actually take that text and turn it into a representation um you need to split it up in a smart way and there's a lot of like thought that goes into how you split that up and I'm not going to go into that cuz I I don't work on that at all but um luckily people have figured out how you can do that intelligently so you split that text up and then then we're going to send all that Downstream so at this point thanks to unstructured and um other capabilities so we had we reached out to Hacker News and we're like what's the latest ID now we have a stream of IDs we send that stream of IDs Downstream uh and now we took those IDs and we enriched them by retrieving metadata from the API which includes included a URL we then requested from that URL the entire HTML we took that HTML and we've processed it with unstructured and now we have we have chunks of text associated with an article uh URL which is associated with an article ID and that article ID is tied to that initial ID that we got from Hecker news um at this point in time we want to take those chunks of text and we want to turn them into a vector and then we'll be able to use the redis connector to write those out to a vector database um and those would be able to be retrieved in our application okay so in order to create an embedding we're going to use a model available on hugging face um and that is also uh going to be quite easy it's actually the easy part of the whole thing because um someone else did all the hard work for us on this one just like instructure did um this is another step in the flow uh wait that's the wrong one so we oh I really got to stop clicking on things all right so we went through um the flow we have the document and now what we're doing is we're looking at the embedding so we're providing a model um pytorch class length of our Vector Dimensions which has to map up with our redus Vector length uh and then the tokenizer so we send that to this this function a little function I wrote but it's essentially we're telling hugging we're we're going to use hugging face Transformers to kind of turn our text our chunk of text into the vector so we'll kind of SK skip over that part a little bit and then we'll get on to the last piece okay so the last part of this um we've gone from IDs to URLs to HTML to chunks of text from chunks of text to vectors now we need to put our vectors into the vector database in order to do that we need to provide the database with some information on what we're going to be storing so we're going to be storing the metadata plus those vectors and we're going to store it in a way so such that we can kind of retrieve potentially all of the uh information associated with the chunk of text because if we think back to the chunked text it's doesn't necessarily represent the entire document and we might want that whole document because there might be a lot of information that we need about that so we'll just store our stuff in redus with a identifier that says hey this is the item ID and then this is the part of it and then in that way we can actually retrieve all the parts so I wrote a little reddis um connector all right so reddis has a library called redis VL which makes it easy to use their Vector I guess it stands for Vector library or something like that I don't know yeah it makes it easier to use the vector uh part of their Vector database uh it's really straightforward there's a similarly to the input um API uh you basically tell the uh worker hey uh we're starting this worker you got to do this first and then each new um batch of data I want you to do that so we've we've written that here we said hey I'm starting that worker um I have a schema go connect to redus let me know if the schema exists if I want you to overwrite it overwrite it and then it's going to run this uh NE B wax will run this when I say it's I mean bite wax will run this build oh really bad poor etiquette here I don't think I can even get back to what I was talking about maybe nope well I actually might have the code somewhere else but hopefully this works no okay so the last part of that was writing the batches so we've connected we've said redus you have this thing don't overwrite it now we have this Handler that says every time we get new data new data in the Stream um that's a batch of data I want you to go out and I'm just going to talk about it and ignore that um we can put this so it's nice to look at uh I want you to go out and do whatever they put in the right batch and in this case it was like go write it into the vector database um so we have the data in the right format we that the schema is um set at and then we're we're good to go so every time we receive a new batch of chunks of text we'll write that batch out um and that's it I had this thing running but since I lost uh I'm gonna it's going to panic but you can see what's happening here so I'll I mean we're asking Hacker News for the IDS we receive IDs we send them Downstream in the code I showed there was an inspect or there was an inspect where I'm logging this out so that's what you're seeing here so basically you know this is a comment so ignore it comment ignore it at some point we received actual story um looks like this is a story here and so we have the text and then we'll have oh no that was a giant comment too so we have there we go so we had a little story it was really short once we pulled it it was looks like a transcript and subtitles uh you can see the text here and then you can see this is the embedding but it's uh bite string that's what reddis required so you can see a bite string list basically here uh and then it failed but this was running in the background when I was like sitting in my booth there so we loaded up a um this Rus database I won't press anything here because I don't want it to um go away on me but uh you can see here that we've shoved this database full of vectors so the next step um that we would do is when we're writing the application to interact with this Vector database where we've created an in like an inest ingestion pipeline running in the background um in this case it's like you know we're going to ask funny questions like this comparison of common list testing Frameworks this is like a a proper little Hacker News piece right here so we would be able to reason about things like that um and especially if they were really trending like I'm sure there was like six or seven uh things about the open AI Dev day uh a week and a half ago or whenever that was um and the way that last piece works is like you'll have to take the request and turn that into an embedding and and then you're going to compare those embeddings in the vector in the vector space um and it will return uh ones that are similar and you can rank them and there's other things you can do there um but that's outside of the scope of this talk uh and then you'll use that in the prompt so this is the document embedding for this one which shows up as funky little bites there and then you can see all the other metadata so in this instance so we see we have key ID so in the if we went back to the code which is probably not in existence anymore um the key ID represents the item so that's like the story and then the chunk so what we could do which is kind of fun is like so the story is 12 chunks long um if you were just using a straight Vector retrieval this might not might not have this optionality but we could actually um do some sort of like Reax search where we could retrieve we could find the top story and then we could retrieve all the chunks because not all the chunks will match potentially but you might want all the chunks and string them together and put that in the context uh yeah that's um that's that's all I got for today um thanks everyone for listening um yeah if you would like to give us a star on on uh GitHub that would be um really helpful thank you I think I have some uh time for questions I think three minutes so if you have any questions I can take them here otherwise um I'll be at the Q&A table and then I'll also be at my booth uh pouring pour over coffee thank you uh I was trying to understand how bite wax is different from another stream processor is it particularly tuned to this embedding sorry I didn't hear the last part it's different is it particularly tuned to the embedding process or is it no so bite wax is just like a generally applicable um framework in processing engine so you can think of it kind of like a Flink or alternative um that's allows you to use Python native code without all the things and you don't it doesn't require you to learn about a a new scheduler or you know task manager or job manager or anything like that yeah just like a quick question uh I noce you use like redice for for the embedding store why not like pine cone or like some you know as a service thing and you decided to go like for like locally deployed R this just curious yeah yeah you you could use whatever you want um you know there's a pine cone python SDK I think or it's just an API call I'm not sure uh so you could use that it would be the exact same pipeline you just swap swap out that last bit so the last line there was um and I you could do quadrant you could use um so this last line here is the output so I wrote a rice Vector output uh class so that's the connector there but you could use whatever you want um we don't have this massive library of connectors but we will in the future um we're gonna have like a connector Hub if you will so you'll be able to like there will be a Rus one a quadrant one elastic search one and then on the input side there's kofka Now red panda and there will be all these different connectors uh We've kept that out of the library for now just to focus on more core problems but there's no so I know Brian who's a developer Advocate at uh redis and I told him like hey I'll do rtis in this one um but I had quadrant before and you could use which really whichever you want good CH makes sense thank yeah of course I think that's it for time indeed thank you thanks everyone [Music]