IndexError: list index out of range for np.argmax()

Question

I have the following code:

import numpy as np
from sklearn.metrics import f1_score
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# Labeled data
data = [
    ("Textbook 1", "This is a great book on natural language processing.", "NLP"),
    ("Textbook 2", "A comprehensive guide to statistical approaches in NLP.", "Statistics"),
    ("Textbook 3", "Learn how to analyze text using the Natural Language Toolkit in Python.", "Python"),
    ("Textbook 4", "An introduction to speech and language processing.", "Speech"),
    ("Textbook 5", "This book covers various NLP techniques and applications.", "NLP")
]

# Split data into features and labels
titles = [d[0] for d in data]
descriptions = [d[1] for d in data]
labels = [d[2] for d in data]

# Initialize TF-IDF vectorizer
vectorizer = TfidfVectorizer()
tfidf_matrix = vectorizer.fit_transform(titles + descriptions)

# Transform query into TF-IDF vector
query = "Textbook on natural language processing techniques with high f1 score evaluation."
query_vector = vectorizer.transform([query])

# Calculate cosine similarity between query and textbooks
cosine_similarities = cosine_similarity(query_vector, tfidf_matrix).flatten()

# Find most similar textbook
most_similar_index = np.argmax(cosine_similarities)
predicted_label = labels[most_similar_index]

# Calculate F1 score
f1 = f1_score([predicted_label], ["NLP"], average="weighted")

# Print predicted label and F1 score
print("Predicted label:", predicted_label)
print("F1 score:", f1)

However I get this error:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[20], line 33
     31 # Find most similar textbook
     32 most_similar_index = np.argmax(cosine_similarities)
---> 33 predicted_label = labels[most_similar_index]
     35 # Calculate F1 score
     36 f1 = f1_score([predicted_label], ["NLP"], average="weighted")

IndexError: list index out of range

The variable 'most_similar_index' returns a value of 5. However I know that for python a list starts at 0 therefore it should be returning 4? How can I fix this?. I think it might have something to np.argmax()?

Answer 1

You have a total of 10 TF-IDF vectors (5 for titles and 5 for descriptions). You can also see the shapes by debugging.

tfidf_matrix = vectorizer.fit_transform(titles + descriptions)

So, it is possible to return 0-10 values. You have to organize it not to get the index out of range error. Pls update your code from line 32:

most_similar_index = np.argmax(cosine_similarities)
if most_similar_index < len(titles):
    predicted_label = labels[most_similar_index]
else:
    predicted_label = labels[most_similar_index - len(titles)]

Answer 2

The + operator between two Python lists doesn't add them element-wise but it creates a new list from extending the first one with the elements of the second one. Therefore, titles + descriptions, and hence cosine_similarities are 10-dimensional arrays.

What you want to do instead is

tfidf_matrix = vectorizer.fit_transform([f"{t} {d}" for t, d in zip(titles, descriptions)])