NER Powered Semantic Search in Python

Key Takeaways

today we're going to go through an example of using named entity recognition alongside a vector search and this is a really interesting way of just making sure our search is very much focused on exactly what ever it is we're looking for so for example if we are going to search through articles and we are searching for something to do with Tesla we can use this to restrict the search whatever it is we're looking for maybe we're looking for some news about Tesla for self-driving and using this we're going to restrict scope for full self driving to specifically articles or possible articles that contain the named entity Tesla now I think this is a really interesting use case and can definitely help making search more specific so let's jump to the example so we're working from this example page on Pine Cone so pine cone IO dots examples ner search and Ashok thought of this and put all of this together yeah I think it's I think it's a really cool example so we'll work through the code so to get started I'm going to come over to colab here and I'm just going to install a few dependencies so it's sentence Transformers python client and data sets okay so that will just install everything okay great and then we can download data that we're going to be using so we're going to go from data sets import load data set we are going to be using the this data set here so it's a medium articles data set I'm going to search for here okay and it's just it contains a ton of Articles script from medium so to use that we use the name of the data set from here and we're going to load it into a hooking face data frame so data FiOS medium articles CSV and we want to translate okay if we have a look at that so we have a look at the head after converting it to a panda's data frame so two pandas as a second moment to download okay and then let's zoom them out a little bit so you can see what we have here so title text you're around a few other things in here okay so obviously the text is going to be the main part for us and what we will also do is just drop a few of these so drop any where we have just empty rows so that's this the drop n a here and then we're going to sample at random 50 000 of these articles and we'll do that with a random state of 32. essentially this is so you can get the same 50 000 articles that I will get here as well okay cool now for each article I mean there's a few things we could do in putting the embeddings together so we're going to have to create embeddings for every single article here what we could do is split the article into Parts because our embedding models can only handle so many words at one time so we could split the article into parts and embed all of those different parts but in this case we can usually get an idea of what the article will talk about based on the title and usually the introduction so what we'll do is actually take the first 1000 characters so what you can see here there and then we're just going to join the article title and the text and we'll keep all of that within the title text feature so we'll go DF head again maybe we can just have a look at title text so let's do that okay so we have this now the next thing to do or that we want to do is initialize our ner model so initialize any r model how do we do that so we're going to be using a hooking face to list so I'll just copy the code in we have all of this so this is a ner Pipeline with an Nar model so we have this D slim base for any R name entity recognition we have our tokenizer the model itself and then we just load all of that into our pipeline we need to select our device so if we have a GPU running we will once you use a GPU it will be much faster so I'm going to import torch and what we're going to do is say device equals Cuda if a coded devices available so torch Dot Cuda is available otherwise we're just going to use CPU now if you are on telev what you can do to make sure you are going to be using GPU is go to runtime change runtime type and here change this to GPU okay so mine wasn't set to GPU tsunami change it save and actually rerun everything so I'll go ahead and do that okay so that all just re-ran and what I'm going to do is just print the device here so you can see that we are in fact using Cuda hopefully okay cool and then here we will run all this and that will just download and initialize the model as well so it may take a moment okay so after that let's try it on a quick example so we have London's cattle England and the United Kingdom so we would expect a few main vertices to be within this so let's run that and see what we return okay with this uh maybe I need to here maybe rewrite torch device this Okay cool so here we have a few things so we have a location this is the entity type London okay cool that is definitely true location again England and location again United Kingdom okay all great so that's definitely working so let's move on to creating our embeddings so to create those embeddings we will need a retrieve model so to initialize that we're going to be using the sentence Transformers library that will look something like this so from sentence Transformers import sentence Transformer and then we initialize this model here so this is just a it's a pretty good sense transform model and it's been trained on a lot of examples so the performance generally really good okay now we can see the model format here so we have match sequence length of 128 tokens word embedded Dimension here 768 so that's how big the sentence embeddings will be sent in selectors and we can use this model to create embeddings but we need someone to start our embeddings and for that we're going to be using pine cones so we will initialize that so for that we will need to do this so import Pinecone initialize our connections Pine counter without an API key which is free and we can get it from here so app.pinecone dot IO we need to log in here now you probably just have one project in here which would probably be your name default project so go into that you go over to API keys and you just press copy on your API key here and then you just paste it into here now I put my API key into a variable called API key and with that we can initialize that connection now we'll create a new index called ner search and we do that with pine cone create index and this index there's a few things that we need here so index we need the index name the dimensionality so Dimension now this is going to be equal to the retriever model that we just created that includes a embedding Dimension and we actually have that up here is 768 but we're going to we're not going to hard code it we're just going to go from here so get it gets get sentence embedding Dimension let me go let me see yeah just here okay so if we just have a look at what that will give us should be 768 yeah and then we also want to use the cosine similarity metric and then that will go ahead and create our own legs after it's been created we'll connect to it so we'll do Pinecone index and we just use it index name again describe the index just to see what is in there which should be nothing Okay cool so we have dimensionality we can see that the index is empty as we haven't added anything into it yet so the next thing is actually preparing all of our data to be added to Pinecone and allowing us to do this search so what we're going to do first is initialize this extract named entities function which is going to extract for a batch of text it's the under abstract name entities for each one those text records so let's initialize that and then what we're going to do is just have a look at how it will work so extract and named entities and in there we're just going to pass the data that we have is called DF title text what we'll do is maybe turn this into a list uh yeah let's do that so DFT text I want to go for the first let's say the first three items okay and we'll turn into a list in our area list there okay cool so the first one we get data Christmas on Miss America light switch so on and so on and yeah you can you can see we're getting those entities they're being extracted from there so that is working and what we can do now is actually Loop through the full data set or the full 50 000 that we have extracted and we're just going to we're going to do this so let me walk walk you through all this so I'm going to be doing everything about Shazam 64. that's to avoid overwhelming our GPU which is going to be encoding everything in batches of 64 and also when we're sending these requests to Pinecone we can't send too many or too larger requests at any one time so with 64 pretty safe here we're going to find the end of the batch so we're going in batch of 64 from I to in obviously when you get to the ended data set the back size will probably be smaller because we probably don't have the data set fitting into perfect 64 batch sizes okay and then we extract the batch it's generate those embeddings let's check the name entities using our function if there are any duplicates so we might have the same name entity a few times in in a single sentence of course or single paragraph in that case we don't need to do that because we're just going to be filtering so actually I need one instance of the named entity so to do that we just deduplicate our entities here and convert it back into a list and then what we're doing here is we're dropping the title text from our batch because we actually don't need that right here because we've just extracted our name empties and that's all we care about at the moment and then from there we create our metadata so this is going to be some metadata that we saw in the dictionary like format in Pinecone create our unique IDs it's just a count and we upset everything or add everything to upset list and upsell so we can run that it will possibly take a little while to run may have to just deal with this quickly so empties can contain a list so probably the best way to deal with that is within the function here so let's come up here and what we'll do is just convert into a or removed duplicates here so any equals less set any okay and then we don't have to do it here so I can remove that and we'll just call this so it's batch named entities Okay cool so let's try and run that again great so that looks like it will work it will take a little bit of time so I will leave that to run and when it's complete we'll go through the actual making queries and see how that works okay so that is done I've just run it on a another fast uh computer so we should be able to describe it in this statue and we should see 50 000 items the one thing here is that I need to just refresh the connection so let me do that quickly so we'll come here I am going to run this again and I'm also going to ruin this without the creating depths come down here run this here we go so we have 50 000 vectors in there now and what we want to do is use this here so this is going to search through Pinecone so we're going to have our query here I'm going to extract the named entities from our query embed our query and what we do is query with our embedded query with our query Vector which is this xq we're going to return the top 10 most similar matches we're going to include metadata for those matches and importantly for the named entity part of this or named entity filtering what we do is we filter for named entities that are within the list that we have here so let's run that this will just return all of the metadata that we have there so the the titles that we that we might want from that and what we'll do is simply make a few queries so what are the best places to visit in Greece will be our first group and we get these titles here so budget friendly holidays was the best summer destination Greece exploring Greece Santorini Island yeah all I think pretty seem like pretty relevant articles to me that's a few more questions what are the best places to visit in London runes and then we get a few more so oh I think pretty relevant again and then let's try one slightly different query or SpaceX don't build a city on Mars and here we go so for sure I think all these are definitely very relevant articles to be returned okay so that's it for this example covering Vector search with this sort of named NT filtering component just to try and improve our results and make sure we are searching within a space that contains whatever it is we are looking for or just making our social scope more specific so I hope this has been interesting and useful thank you very much for watching and I will see you again in the next one bye