Spaces:
Sleeping
Sleeping
File size: 3,520 Bytes
ff7b988 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | import os
import re
accepted = frozenset([chr(i) for i in range(ord('a'), ord('z') + 1)] +
[chr(i) for i in range(ord('A'), ord('Z') + 1)] +
[chr(i) for i in range(ord('0'), ord('9') + 1)])
rex = re.compile(r'_+')
def norm(s):
s = ''.join([c if c in accepted else '_' for c in s.lower()])
s = rex.sub('_', s).strip('_')
return s
def create_reverse_lookup(atoi):
# Multiple entries could go to the same artist_id/genre_id
itoa = {}
for a, i in atoi.items():
if i not in itoa:
itoa[i] = []
itoa[i].append(a)
indices = sorted(list(itoa.keys()))
for i in indices:
itoa[i] = '_'.join(sorted(itoa[i]))
return itoa
class ArtistGenreProcessor():
def __init__(self, v3=False):
self.v3 = v3
dirname = os.path.dirname(__file__)
if self.v3:
self.artist_id_file = f"{dirname}/ids/v3_artist_ids.txt"
self.genre_id_file = f"{dirname}/ids/v3_genre_ids.txt"
else:
self.artist_id_file = f"{dirname}/ids/v2_artist_ids.txt"
self.genre_id_file = f"{dirname}/ids/v2_genre_ids.txt"
self.load_artists()
self.load_genres()
def get_artist_id(self, artist):
input_artist = artist
if self.v3:
artist = artist.lower()
else:
artist = norm(artist)
if artist not in self.artist_ids:
print(f"Input artist {input_artist} maps to {artist}, which is not present in {self.artist_id_file}. "
f"Defaulting to (artist_id, artist) = (0, unknown), if that seems wrong please format artist correctly")
return self.artist_ids.get(artist, 0)
def get_genre_ids(self, genre):
if self.v3:
genres = [genre.lower()]
else:
# In v2, we convert genre into a bag of words
genres = norm(genre).split("_")
for word in genres:
if word not in self.genre_ids:
print(f"Input genre {genre} maps to the list {genres}. {word} is not present in {self.genre_id_file}. "
f"Defaulting to (word_id, word) = (0, unknown), if that seems wrong please format genre correctly")
return [self.genre_ids.get(word, 0) for word in genres]
# get_artist/genre throw error if we ask for non-present values
def get_artist(self, artist_id):
return self.artists[artist_id]
def get_genre(self, genre_ids):
if self.v3:
assert len(genre_ids) == 1
genre = self.genres[genre_ids[0]]
else:
genre = '_'.join([self.genres[genre_id] for genre_id in genre_ids if genre_id >= 0])
return genre
def load_artists(self):
print(f'Loading artist IDs from {self.artist_id_file}')
self.artist_ids = {}
with open(self.artist_id_file, 'r', encoding="utf-8") as f:
for line in f:
artist, artist_id = line.strip().split(';')
self.artist_ids[artist.lower()] = int(artist_id)
self.artists = create_reverse_lookup(self.artist_ids)
def load_genres(self):
print(f'Loading artist IDs from {self.genre_id_file}')
self.genre_ids = {}
with open(self.genre_id_file, 'r', encoding="utf-8") as f:
for line in f:
genre, genre_id = line.strip().split(';')
self.genre_ids[genre.lower()] = int(genre_id)
self.genres = create_reverse_lookup(self.genre_ids)
|