Views
No views yet
1 #! pip install transformers
2 #! pip install torch
3 #! pip install datasets1from transformers import DistilBertForSequenceClassification, DistilBertTokenizer
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3from datasets import load_dataset
4import numpy as np
5import torch
6
7
8dataset = load_dataset("manoh2f2/songs_resampled")
9
10# Cargar el dataset en un DataFrame
11split_name = 'train'
12df_resampled = dataset[split_name].to_pandas()
13
14tokenizer = AutoTokenizer.from_pretrained("manoh2f2/recommend_songs")
15model = AutoModelForSequenceClassification.from_pretrained("manoh2f2/recommend_songs")
16
17# Define a prompt
18prompt = "I am happy"
19
20# Tokenize the prompt
21encoded_prompt = tokenizer(prompt, return_tensors='pt', max_length=256)
22
23# Make a prediction using the trained model
24with torch.no_grad():
25 model_output = model(**encoded_prompt)
26
27# Get the predicted emotion index
28predicted_emotion_index = torch.argmax(model_output.logits).item()
29
30# Map the index back to the emotion label using the DataFrame
31predicted_emotion_label = df_resampled['emotions'].unique()[predicted_emotion_index]
32
33# Get a song associated with the predicted emotion from the DaraFrame
34result = df_resampled[df_resampled['emotions'] == predicted_emotion_label]
35
36# Get the number of rows in the DataFrame
37num_rows = result.shape[0]
38#Generate a random index to select a random song from the DataFrame
39random_index = np.random.randint(0, num_rows)
40
41#Get the recommended song and artist
42recommended_song = result['song'].iloc[random_index]
43recommended_artist = result['artist'].iloc[random_index]
44
45#Print the results
46print(f"Prompt: {prompt}")
47print(f"Predicted Emotion: {predicted_emotion_label}")
48print(f"Recommended Song: {recommended_song} - {recommended_artist}")