ModusMusic / song_matching.py
IshmamF's picture
Update song_matching.py
d99bf7c
raw
history blame
2.05 kB
import pandas as pd
from sentence_transformers import SentenceTransformer, util
class SongMatcher:
def __init__(self, songs_data_file, model_name="sentence-transformers/all-mpnet-base-v2"):
"""
Initializes the SongMatcher with the songs data file and the SentenceTransformer model.
:param songs_data_file: Path to the CSV file containing songs data
:param model_name: Name of the SentenceTransformer model
"""
self.songs_df = pd.read_csv(songs_data_file)
self.sim_model = SentenceTransformer(model_name)
def match_songs_with_sentiment(user_sentiment_label, user_sentiment_score, user_input, score_range):
# New: Define inputVector here
inputVector = self.sim_model.encode(user_input)
# Filter songs with the same sentiment label
matched_songs = self.songs_df[self.songs_df['sentiment'] == user_sentiment_label]
# Calculate the score range
score_min = max(0, user_sentiment_score - score_range)
score_max = min(1, user_sentiment_score + score_range)
# Further filter songs whose scores fall within the specified range
matched_songs = matched_songs[(matched_songs['score'] >= score_min) & (matched_songs['score'] <= score_max)]
# Shuffle the matched songs to get a random order
matched_songs = matched_songs.sample(frac=1).reset_index(drop=True)
matched_songs['similarity'] = matched_songs['seq'].apply(lambda x: util.pytorch_cos_sim(sim_model.encode(x), inputVector))
top_5 = matched_songs['similarity'].sort_values(ascending=False).head(5)
# Sort the songs by how close their score is to the user's sentiment score
# matched_songs['score_diff'] = abs(matched_songs['score'] - user_sentiment_score)
# matched_songs = matched_songs.sort_values(by='score_diff')
# Select the top five songs and return
return matched_songs.loc[top_5.index, ['song','artist','seq','similarity','sentiment','score']]