Spaces:
Sleeping
Sleeping
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(self, user_sentiment_label, user_sentiment_score, user_input, score_range=0.00625): | |
""" | |
Matches songs from the dataset with the user's sentiment. | |
:param user_sentiment_label: The sentiment label of the user input | |
:param user_sentiment_score: The sentiment score of the user input | |
:param user_input: Text input from the user | |
:param score_range: Range for filtering songs based on sentiment score | |
:return: DataFrame of top 5 matched songs | |
""" | |
# 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)] | |
# Compute similarity between user input and song lyrics | |
input_vector = self.sim_model.encode(user_input) | |
matched_songs['similarity'] = matched_songs['seq'].apply(lambda x: util.pytorch_cos_sim(self.sim_model.encode(x), input_vector)) | |
# Select the top five songs based on similarity and return | |
top_5 = matched_songs.nlargest(5, 'similarity') | |
return top_5[['song', 'artist', 'seq', 'similarity', 'sentiment', 'score']] | |