Semantic Textual Search with Vector Embeddings

!pip install -qU pinecone-client ipywidgets

import pinecone

# Load Pinecone API key
import os
api_key = os.getenv("PINECONE_API_KEY") or "YOUR-API-KEY"
pinecone.init(api_key=api_key)
# List all indexes currently present for your key
pinecone.list_indexes()

!pip install -qU wordcloud pandas-profiling
!pip install -qU sentence-transformers --no-cache-dir

import pandas as pd
import numpy as np
import time
import re
from wordcloud import WordCloud, STOPWORDS, ImageColorGenerator
import matplotlib.pyplot as plt
import sqlite3

pd.set_option('display.max_colwidth', 200)

# Pick a name for the new index
index_name = 'semantic-text-search'

# Check whether the index with the same name already exists
if index_name in pinecone.list_indexes():
    pinecone.delete_index(index_name)

pinecone.create_index(name=index_name, metric='cosine', shards=1)

index = pinecone.Index(name = index_name, response_timeout=300)
index.info()

import requests, os
DATA_DIR = 'tmp'
URL = "https://www.dropbox.com/s/b2cyb85ib17s7zo/all-the-news.db?dl=1"
FILE = f"{DATA_DIR}/all-the-news.db"

def download_data():
    os.makedirs(DATA_DIR, exist_ok=True)

    if not os.path.exists(FILE):
        r = requests.get(URL)  # create HTTP response object
        with open(FILE, "wb") as f:            
            f.write(r.content)

download_data()

cnx = sqlite3.connect(FILE)
data = pd.read_sql_query("SELECT * FROM longform", cnx)
data.set_index('id', inplace=True)
data.head()

# Define number of test articles
NUM_OF_TEST_ARTICLES = 2

# Remove test articles from data and keep them in separate dataframe
test_articles = data[['title','content']][97::81][:NUM_OF_TEST_ARTICLES]
data.drop(list(test_articles.index), inplace=True)

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('average_word_embeddings_komninos')

# Fill missing and remove redundant data
data['title'] = data['title'].fillna('')

# Create vector embeddings based on the title column
print('Encoding titles...')
encoded_titles = model.encode(data['title'].tolist(), show_progress_bar=True)
data['title_vector'] = encoded_titles.tolist()

# Upsert title vectors in title namespace
print("Uploading vectors to title namespace..")
acks_titles = index.upsert(items=zip(data.index, data.title_vector), namespace='title', batch_size=1000)

# Fill missing data
data['content'] = data['content'].fillna('')

# Extract only first few sentences of each article for quicker vector calculations
data['content'] = data.content.apply(lambda x: ' '.join(re.split(r'(?<=[.:;])\s', x)[:10]))

# Create vector embeddings based on the content column
print('Encoding content...')
encoded_content = model.encode(data['content'].tolist(), show_progress_bar=True)
data['content_vector'] = encoded_content.tolist()

# Upsert content vectors in content namespace
acks_content =index.upsert(items=zip(data.index, data.content_vector), namespace='content', batch_size=1000)

# Check index size for each namespace
print(index.info(namespace='title'))
print(index.info(namespace='content'))

# Print test articles
display(test_articles)

titles_mapped = dict(zip(data.index, data.title))
content_mapped = dict(zip(data.index, data.content))

def get_wordcloud_for_article(recommendations, namespace):
    'Generates word cloud for the recommendations (titles or content).'

    stopwords = set(STOPWORDS).union([np.nan, 'NaN', 'S'])
    wordcloud = WordCloud(
                   max_words=50000, 
                   min_font_size =12, 
                   max_font_size=50, 
                   relative_scaling = 0.9, 
                   stopwords=set(STOPWORDS),
                   normalize_plurals= True
                  )
    
    if namespace == 'title':
        clean_titles = [word for word in recommendations.title.values if word not in stopwords]
        wordcloud = wordcloud.generate(' '.join(clean_titles))
    else:
        clean_content = [word for word in recommendations.content.values if word not in stopwords]
        wordcloud = wordcloud.generate(' '.join(clean_content))

    plt.imshow(wordcloud, interpolation='bilinear')
    plt.axis("off")
    plt.show()


def print_query_results(query_result, query, namespace, show_options={'wordcloud':True, 'tabular':True}):
    'Prints query result with wordcloud.'
      
    print(f'\nMost similar results querying {query} in "{namespace}" namespace:\n')

    res = query_result[0]
    df = pd.DataFrame({'id':res.ids, 
                       'score':res.scores,
                       'title': [titles_mapped[int(_id)] if int(_id) in titles_mapped else ' '  for _id in res.ids],
                       'content': [content_mapped[int(_id)] if int(_id) in content_mapped else ' '  for _id in res.ids],
                       })
    if show_options['tabular']:
        display(df.head(5))
    if show_options['wordcloud']:
        get_wordcloud_for_article(df, namespace)
    print('\n')

def query_article_title(test_article, namespace, top_k=5, show_options={'wordcloud':True, 'tabular':True}):
    '''Queries an article using its title in the specified
     namespace and prints results.'''

    # Create vector embeddings based on the title column
    encoded_titles = model.encode(test_article['title'], 
                                  show_progress_bar=False)
    test_article['title_vector'] = pd.Series(encoded_titles.tolist())

    # Query namespace passed as parameter using title vector
    query_result_titles = index.query(queries=[test_article.title_vector], 
                                      namespace=namespace, 
                                      top_k=top_k, 
                                      disable_progress_bar=True)

    # Print query results 
    if show_options['wordcloud'] or show_options['tabular']:
        print_query_results(query_result_titles, query='title', 
                            namespace=namespace, 
                            show_options=show_options)

    return query_result_titles

def query_article_content(test_article, namespace, top_k=5, show_options={'wordcloud':True, 'tabular':True}):
    '''Queries an article using its content in the specified 
    namespace and prints results.'''

    # Create vector embeddings based on the content column
    encoded_content = model.encode(test_article['content'], 
                                   show_progress_bar=False)
    test_article['content_vector'] = pd.Series(encoded_content.tolist())

    # Query content namespace using content vector
    query_result_content = index.query(queries=[test_article.content_vector], 
                                       namespace=namespace, 
                                       top_k=top_k, 
                                       disable_progress_bar=True)

    # Print query results 
    if show_options['wordcloud'] or show_options['tabular']:
        print_query_results(query_result_content, 
                            query='content', 
                            namespace=namespace, 
                            show_options=show_options)

    return query_result_content

def aggregate_results(article):
    '''Aggregates results after querying both namespaces
       for both the article's title and content.'''

    results = []
    
    results.append(query_article_title(article, namespace='title', top_k=30, show_options={'wordcloud':False, 'tabular':False}))
    results.append(query_article_title(article, namespace='content', top_k=30, show_options={'wordcloud':False, 'tabular':False}))
    results.append(query_article_content(article, namespace='title', top_k=30, show_options={'wordcloud':False, 'tabular':False}))
    results.append(query_article_content(article, namespace='content', top_k=30, show_options={'wordcloud':False, 'tabular':False}))

    articles_scores = {}
    articles_count = {}

    for res in results:
        for id, score in zip(res[0].ids, res[0].scores):
            if id not in articles_scores:
                articles_scores[id] = score
                articles_count[id] = 1
            else:
                articles_scores[id] += score
                articles_count[id] += 1
    
    return articles_scores , articles_count

def show_aggregated_results(results_dict, counts, show_options={'wordcloud':True, 'tabular':True}):
    '''Shows results after aggregation. Values are sorted based
    on the number of queries they appear (1-4) and based on their
    average score.'''
    
    df = pd.DataFrame({'id':results_dict.keys(), 
                       'count': counts.values(),
                       'average_score':[round(r/c, 3) for r, c in zip(results_dict.values(),counts.values())],
                       'title': [titles_mapped[int(_id)] if int(_id) in titles_mapped else ' '  for _id in results_dict.keys()],
                       'content': [content_mapped[int(_id)] if int(_id) in content_mapped else ' '  for _id in results_dict.keys()],
                       })
    df.sort_values(by=['count', 'average_score'], ascending=False, inplace=True)
    
    if show_options['tabular']:
        print('\nMost similar results after aggregation:\n')
        display(df.head(5))
    if show_options['wordcloud']:
        print('\nWordcloud for titles and content after aggregation:')
        print('-Titles:')
        get_wordcloud_for_article(df[:10], 'title')
        print('-Content:')
        get_wordcloud_for_article(df[:10], 'content')
    print('\n')

# Query index using simple and cross namespace approach
for e, (_, test_article) in enumerate(test_articles.iterrows()):
    print(f'\nArticle {e+1}')
    print(f'\n Title: {test_article.title}')
    print(f' Content: {test_article.content[:200].strip()}' + ('...' if len(test_article.content) > 200 else ''))
    
    # Uncomment to query the titles in title namespace
    #query_article_title(test_article, 'title',  show_options={'wordcloud':True, 'tabular':False})

    # Uncomment to query the content in content namespace
    #query_article_content(test_article, namespace='content', show_options={'wordcloud':True, 'tabular':False})

    # Cross namespace query
    aggregated_results, counts = aggregate_results(test_article)
    show_aggregated_results(aggregated_results, counts, show_options={'wordcloud':False, 'tabular':True})

pinecone.delete_index(index_name)