Fine-Tuning Text Embeddings For Domain-specific Search (w/ Python)

Key Takeaways

embedding models represent text as semantically meaningful vectors although they can be readily used for countless use cases like retrieval or classification general purpose embedding models may have poor performance on domain specific tasks in this video I'll discuss how we can overcome this limitation via fine tuning I'll start with a highlevel overview of key Concepts and then walk through a concrete example of fine-tuning embeddings on AI J posts okay so let's talk about fine-tuning text embeddings this video is going to assume you already have a basic understanding of what text embeddings do but if you need a refresher I gave a beginner friendly introduction in a previous video so you can check that out if you need to so just jumping right into it Tex embeddings have become more popular these days with the widespread usage of so-called retrieval augmented generation or rag for short and the basic idea behind rag is to improve llm based systems by retrieving relevant content whenever passing a prompt to a large language model a basic example of this is we might have a FAQ bot or a question answering AI assistant where we can pass in a query instead of just passing this query directly to a large language model in a rag system there will be this pre-step of context retrieval and then that query and the context are combined into the prompt which is then fit the large language model to generate a response where embeddings come into this is typically in this context retrieval step the way that works is you'll create this thing called a vector database which has kind of become a buzzword more recently the basic idea of this is taking a set of items in a knowledge base this is typically like chunks of PDF documents or something like that and instead of representing them in their original text in a vector database they are represented by semantic basally meaningful vectors in other words text embedding as a pre-step to building out this retrieval process you'll create this Vector database so at inference time when a query comes in you can simply embed it and then do this so-called similarity search or semantic search where you compute the similarity between the vector representation of the query and all the different items in your knowledge base or your vector database although doing retrieval in this way has become very popular in rag system systems there's one key problem with it and it boils down to this similar doesn't necessarily mean helpful so let me break this down a little bit here in similarity search or also called semantic search we have a vector representation of a query and we have Vector representations of items in a knowledge base we can compute search results by Computing the similarity between this Vector representation of the query and all the items in the knowledge base to make this more concrete the query might be something like how do I update my payment method this might be a very common question so we would expect that it's pretty easy for our system to answer something like this however in practice one thing that might happen is when we do this similarity search the results might be something like this where the top result so in other words the most similar chunk is something like to view your payment history visit the billing section of your account the second most similar result might be payments can be made via credit cards debit cards a and mobile payments it's only the third most similar result that we get something that seems helpful which is to update billing details visit the billing section of your account so this is kind of like the central issue with similarity search which is that just because a query and an item in your knowledge base are similar that doesn't necessarily mean that the item is helpful in answering the query so one thing we can do to overcome these issues is to fine-tune and embedding model for a specific use case so the basic idea of fine-tuning is to adapt a model to a particular use case through additional training so the intuition here is that we might have an original base model which is like a blank slab of marble and the fine tuning process is basically just refining and chiseling this slab into something that is more appropriate for our use case and then of course there's nothing special about a fine-tuned model and it's common practice to fine-tune models multiple times to give it different behaviors and so we we could do additional fine tuning to take a fine-tuned model and make it even more optimized for our use case and so I've actually made a handful of videos on fine-tuning if this is a new idea to you I would recommend checking those out but specifically when it comes to text embeddings the central motivation for fine-tuning is that this idea of similarity varies by use case for example consider these two pieces of text we have one that says what is Rag and another one that says what is fine tun so let's say we were trying to train a comment classifier to basically take YouTube comments and cluster similar ones together and ensure that different ones aren't clustered together in that case these two questions should be represented in similar Ways by our embedding model so that when we do the clustering they will be put into the same group however if we were doing something like an FAQ search the questions what is Rag and what is fine tuning these should be dissimilar things because if we have an incoming query that says what is rag an item in the knowledge base that contains the text what is fine-tuning wouldn't be helpful in answering this query so ideally these should be dissimilar and taking this one step further if we were talking about a different context so instead of doing FAQ search in the context of AI let's say we're doing document search in the context of project management where rag instead of standing for retrieval augmented generation it instead stands for red Amber and green which is like this stoplight system used sometimes in project management then these items should again be dissimilar because this question isn't helpful in answering this query but the representations should be different because the meaning of rag in the context of project management is much different than its meaning in AI so here I'll break down the embedding fine-tuning process into five basic steps the first step is to get rather positive and negative pairs the second step is to pick a pre-trained model and we'll see how we can do this the third step is to pick a loss function which will basically guide the fine-tuning process and then this will depend on the task that you're trying to do and the data that you have available the fourth step is you're going to fine-tune the pre-trained model using your data set and then finally we're going to evaluate the model so instead of talking about these five steps in an abstract way I think it would be helpful to walk through a concrete example so we can see what each of these steps looks like in code so here I'm going to fine-tune a set of text embeddings on AI job listings the basic idea is here we want to take humanlike queries like data scientist 6 years of experience llms credit risk content marketing and we want to match this query with some kind of job post there are two things here that make fine-tuning helpful for this task one is that the the queries and the job descriptions look very different queries are typically very short while job descriptions can be quite long there's a need right there for some kind of fine tuning and additionally there's a lot of technical jargon in the AI space that may not be captured well by just general purpose embedding models because of these two different aspects of this problem fine-tuning can be a good solution starting with that first step of gathering positive and negative pairs to make this a bit more concrete we might have a queer that is something like data pipeline management Advanced SQL techniques ETL elt processes a positive match might be a job description that has something like this make your impact within a rapidly growing fintech company build and manage robust data pipelines that support scalable and then it goes on and on and on basically this is a good match cuz it seems like this role is for a data engineer and this query is also for a data engineer so it seems like this is a good match a negative match on the other hand might look something like this collaborate to design and imp Implement endtoend Ai and data science life cycles from data collection and pre-processing to model deployment and model monitoring so this job on the other hand sounds more like a data scientist or ml engineering role which isn't so well suited for this query so basically what we want to do is gather a bunch of these examples so we can train an embedding model to be more aligned with the behavior that we want it to have while this might sound simple enough this first step is by far the most important and timeconsuming part of this entire fine-tuning process just because you got to ask okay where am I going to get this data from and how am I going to establish positive matches how I'm going to do the negative matches so on and so forth so it may not be obvious how to generate this data but lucky for us I've already gone through the trouble of doing that and I've uploaded the data set on the hugging face Hub so we can just import it with one line of code so I won't walk through all of the code I use to do this because it's quite long and we would spend most of the video talking about it but I will kind of walk walk through the process at a high level and then call out which notebooks and scripts you can check out that correspond to each of these steps and actually this is a more streamlined process than the original version because originally I had gone through a process of scraping a job board to get like a thousand different AI jobs and then go through this whole process but then after I had already uploaded the data and trained the model and all that I realized that the websites terms of use said you can't train AIS on their data so I basically had to start over this is actually a more streamlined process because I just extracted the AI job descriptions from a hugging face data set rather than scraping a website on my own and I guess that's just a good lesson of before you go through the trouble of scraping a website and doing a whole project on it check out the terms of use of that website and then they typically have a robots. text file that you can check out that lets you know what pages are okay or not okay to scrap but just walking through this process step by step the first thing I did was extract the AI job descriptions just imported it from this data set on the hugging face hub from data Stacks called LinkedIn job listing so you can also grab this data set if you like it has 120,000 jobs on it the second thing I did to generate the humanik queries I actually use GPT 40 mini to do this so basically I took the job descriptions from this LinkedIn job listing data set I crafted this prompt for gbt 40 mini that basically said something like given this job description generate a concise like query that corresponds to this description then finally I took all those prompts and then I used open ai's batch API to generate queries for all the job descriptions so if you want to see how I did that you can check out this notebook here it's called one generate synthetic queries. II andb and that's available on the GitHub repository linked here now that we have the humanik queries and we have the job descriptions we'll need to create our positive pairs so basically I took the job descriptions next clean them up and this is a key point so I removed parts that were basically irrelevant to the qualifications of the job boilerplate text that is typically in job descriptions like about the company or equal employment opportunities and things like that I then matched up these cleaned job descriptions to the synthetic queries and then I stored them all in a pandas data frame this is what the result looks like and we could actually stop here because there are fine-tuning approaches that only require positive pairs they don't require negative pairs but I thought to just go one step further and to also get the negative pairs so we can see what that looks like to generate the negative pairs I did something inspired by the generative pseudo labeling paper which I'll pop on the screen here where I basically took all the job descriptions I computed the similarities between them the result of that is this Square Matrix consisting of similarity scores between all the job descriptions and then what I did was for each query I matched the least similar job description as the negative pair the negative example while ensuring there were no repetitions if you're curious about how I did that you can check out the create training data. II Notebook available on the GitHub step two now that we have our data is to pick a pre-trained model and so to do this I went to the sentence Transformers Library so esper.net and there they have these various leaderboards ranking different embedding models so if you go to reference 5 Linked In the description you'll see various base models and their performance on 14 different sentence similarity tasks also if you check out reference six they have various embedding models fine tuned on semantic search specifically so matching queries to relevant items and basically what I did to pick a pre-train model is I passed a handful of different options through this triplet evaluator so this is just a handy object in the sentence Transformers library that allows you to take a data set and create this evaluator with it which you can can pass in a model to and it'll compute the accuracy so I tried just like a few different models nothing too rigorous and found that this one here the all distill Roberta V1 had the best accuracy of about 88% and so I went ahead with this one if you're doing a different kind of task this may not be the best model for you so it's worth reviewing the different options in the sentence Transformers library and seeing what works best on your specific data and then I will point out I'm using that data set I created specifically for this example which has a train test and validation split so here I'm using the validation data to pick the pre-trained model as opposed to the training data or the testing data and that's important because validation data is typically what you use to tune hyperparameters now that we have our data and pre-train model selected next we can select a loss function here I used this handy summary table in the sentence Transformers doc so this is really helpful because they break down what loss functions are appropriate for a given type of data so for example if you just have single sentences and a class so this is some kind of like text classification task they have various loss functions that make sense for that if you're just trying to do domain adaptation on single sentences without a particular label you can use these loss functions they've anchor anchor pairs which not entirely sure what differentiates this from anchor positive pairs but maybe someone will tell me in the comments but for our use case of doing search we're going to be interested in these so anchor would be something like our query and then we have a positive match so we see we have several loss functions that make sense here we might have data structured such that we have an anchor and a positive or negative match and then we have a label one or zero if it's negative so this just like a slightly different way to structure the data and then finally we have this triplet which is what I've talked about here so far so instead of having a flag to indicate whether it's positive or negative you can have the query the positive match and the negative match all grouped together in this triplet here I have data structured like this so we have a few different options and I just went ahead and used this multiple negatives ranking loss they had some example code on their documentation using it so I feel like if they're going to use it for example code it's probably pretty good choice but you know might be worth experimenting with other loss functions and seeing how those work out on this specific use case now that we have our data our model in our loss function we can start training the model we first need to define various hyperparameters so here I followed closely the example code in sentence Transformers documentation just making some small tweaks one important note for contrastive learning which is what this type of machine learning is called where you're training on these positive and negative pairs is you want to have big batch sizes so I was just training this on my laptop so I didn't want to go too crazy so I did a batch size of 16 I I feel like I could have gone bigger actually cuz I just got a new laptop that has 64 GB of memory maybe I could have tried 32 or 64 but this seemed to be a decent starting point at least and then there are various training arguments we can set here so we can Define where we want the model to be saved after training the number of epoch I just went with one the train and evaluation batch size which I just went with 16 learning rate 2 to Theus 5 got this warm-up ratio did this batch sampler no duplicate this was just from sentence Transformers example code basically it ensures that there are no duplicates in each batch because they are randomly sampled and apparently the specific loss function that we used here benefits from no duplicates in the batch eval strategy is basically how often we want to evaluate our loss here I'm going to do it every 100 steps going to compute the loss of the validation data set and then print it out every 100 steps the data set here is about a th000 or I guess 800 training examples so they should print out eight different losses during training those are all our arguments so we can just pass all these into a sentence Transformers trainer which takes things like the model the training arguments the data that we want to use the loss function and the evaluator so this is the same thing I used at the beginning to evaluate the different pre-trained models it's that same object so this took about 45 seconds to run on my new Macbook which is in St contrast to my M1 Mac Mini which this would have taken two 2 hours to run but anyway we can jump to step five and evaluate the model this is just a screenshot from The Notebook on GitHub so I created another one of these triple evaluators but now for the testing data set rather than the validation data and then I just print the accuracy on the validation data set and the testing data set so we can see that the validation performance jumped from 88% to 99% and then we see that the testing performance is 100% this might indicate really good results it also might indicate that the data that we're using is somewhat artificial which is one of the reasons why I would have preferred to use the raw real data I scraped from that website but I think this example still gets the point across of like how to do this whole process even if we are grabbing a data set that's readily available then as a bonus you can use the model so I pushed the model to the hugging face Hub so anyone can use it and then if you want to use it for inference this is the syntax so you can import it using sentence Transformers Define a new query then you can encode the query like this to generate embeddings and then you can take whatever job descriptions you like encode them and compute the similarities between the query and all the different job descriptions okay so the main reason I made this video was to actually set up the next video that I want to make which is fine-tuning multimodal embeddings so I put out a poll on YouTube and Linkedin for video ideas and this was the most popular on both platforms so I thought it would be helpful to the first do fine-tuning text embeddings to set us up for this next one here conceptually we're going to do basically the same thing in the multimodal setting but instead of just operating on text Data we're going to operate on multiple types of data modalities and so one situation where this might be helpful is on YouTube where you have text that comes from titles and then you have images that come from thumbnails and let's say you want to do some kind of thumbnail Search tool or title search tool or you have a really good title idea and you're trying to match it to a thumbnail idea multimodal embeddings would be something you could use for that use case but the problem is if I were to pass this title and this thumbnail into a pre-trained multimodal model like clip we might get low similarity because there's no reason for it to think that these are similar however through fine-tuning so specifically fine-tuning clip on all my titles and all my thumbnails we can adapt this model to say that these two things have high similarity so in the future if I have a title idea or a thumbnail idea I could generate embeddings for it and then try to do some kind of search over a thumbnail library or title library to help in the idea generation process so that's planned for next time and this video and the next video are going to be part of this new fine-tuning playlist that I put together I've realized I've made a handful of videos on fine tuning and I see myself making a handful more this year with my new laptop so I wanted to have a central place for all of the fine-tuning content fine tuning is fun for me and also thinking about product development model fine tuning is an easy way to differentiate your product in a world full of GPT rappers because at the end of the day your product uses a fine-tuned model the value that you're providing is in that data set that you use to fine-tune it rather than the off-the-shelf model that literally everyone in the world has access to so if you're interested in fine-tuning this playlist will be a great resource and if you have any suggestions for future content to cover in this series please let me know in the comment section below and as always thank you so much for your time and thanks for watching