Update app.py

app.py

@@ -1,19 +1,21 @@
 from sentence_transformers import SentenceTransformer, util
 from huggingface_hub import hf_hub_download
+import os
 import pickle
 import pandas as pd
 import gradio as gr
 
 pd.options.mode.chained_assignment = None  # Turn off SettingWithCopyWarning
 
-pickled = pickle.load(open(hf_hub_download("NimaBoscarino/playlist-generator", filename="clean-large_embeddings_msmarco-MiniLM-L-6-v3.pkl"), "rb"))
-songs = pd.read_csv(hf_hub_download("NimaBoscarino/playlist-generator", filename="songs_new.csv"))
-verses = pickle.load(open(hf_hub_download("NimaBoscarino/playlist-generator", filename="verses.pkl"), "rb"))
-lyrics = pd.read_csv(hf_hub_download("NimaBoscarino/playlist-generator", filename="lyrics_new.csv"))
+auth_token = os.environ.get("TOKEN_FROM_SECRET") or True
+pickled = pickle.load(open(hf_hub_download("NimaBoscarino/playlist-generator", repo_type="dataset", filename="clean-large_embeddings_msmarco-MiniLM-L-6-v3.pkl"), "rb"))
+songs = pd.read_csv(hf_hub_download("NimaBoscarino/playlist-generator", repo_type="dataset", filename="songs_new.csv"))
+verses = pd.read_csv(hf_hub_download("NimaBoscarino/playlist-generator-private", repo_type="dataset", filename="verses.csv", use_auth_token=True))
+lyrics = pd.read_csv(hf_hub_download("NimaBoscarino/playlist-generator-private", repo_type="dataset", filename="lyrics_new.csv", use_auth_token=True))
 
 embedder = SentenceTransformer('msmarco-MiniLM-L-6-v3')
 
-genius_ids = pickled["genius_ids"]
+song_ids = pickled["song_ids"]
 corpus_embeddings = pickled["embeddings"]
 
 
@@ -23,10 +25,10 @@ def generate_playlist(prompt):
     hits = pd.DataFrame(hits[0], columns=['corpus_id', 'score'])
 
     verse_match = verses.iloc[hits['corpus_id']]
-    verse_match = verse_match.drop_duplicates(subset=["genius_id"])
-    song_match = songs[songs["genius_id"].isin(verse_match["genius_id"].values)]
-    song_match.genius_id = pd.Categorical(song_match.genius_id, categories=verse_match["genius_id"].values)
-    song_match = song_match.sort_values("genius_id")
+    verse_match = verse_match.drop_duplicates(subset=["song_id"])
+    song_match = songs[songs["song_id"].isin(verse_match["song_id"].values)]
+    song_match.song_id = pd.Categorical(song_match.song_id, categories=verse_match["song_id"].values)
+    song_match = song_match.sort_values("song_id")
     song_match = song_match[0:9]  # Only grab the top 9
 
     song_names = list(song_match["full_title"])
@@ -40,7 +42,7 @@ def generate_playlist(prompt):
 
 
 def set_lyrics(full_title):
-    lyrics_text = lyrics[lyrics["genius_id"].isin(songs[songs["full_title"] == full_title]["genius_id"])]["text"].iloc[0]
+    lyrics_text = lyrics[lyrics["song_id"].isin(songs[songs["full_title"] == full_title]["song_id"])]["text"].iloc[0]
     return gr.Textbox.update(value=lyrics_text)