Search code examples
python-3.xpandaslistdataframesklearn-pandas

Trying to return the row if sentence is present in pandas dataframe with index value


I have one dataframe. I'm implementing sentence transformers and returning one row based on the search query. For example

search_string = "thor's weapon"
search_vect = model.encode([search_string])
K = 3 # no. of paragraphs that has to be extracted
distilbert_similar_indexes = find_similar(search_vect, embeddings_distilbert, K)
print(distilbert_similar_indexes)

The output is :

array([7, 6, 1], dtype=int64)

Code followed by :

output_data = []
for index in distilbert_similar_indexes:
    output_data.append(paragraph[index])
print(output_data[0])

The output is :

'Stormbreaker is an enchanted axe used by Thor. It was forged from Uru on Nidavellir, and can summon the Bifrost.'

Now i want to search the above sentence from the given dataframe and want to return the single row with that sentence and all given sentences. For example if i have dataframe like

Sentences                                                Tag
0   Thor Odinson is the Asgardian God of Thunder, ...    Excel 1
1   Upon being welcomed back to Asgard as a hero, ...    Excel 1
2   Thor returned to Asgard having defeated his br...    Excel 2
3   Loki Laufeyson was the biological son of Laufe...    Excel 1
4   Stormbreaker is an enchanted axe used by Thor....    Excel 3

Now i want to return the below row

4   Stormbreaker is an enchanted axe used by Thor....    Excel 3

Solution

  • You could try this:

    slice = df[df["Sentences"].str.contains(output_data[0])]
    
    print(slice)
    # Outputs
    4   Stormbreaker is an enchanted axe used by Thor....    Excel 3