Spaces:
Sleeping
Sleeping
| 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) | |