Update app.py

app.py

@@ -6,27 +6,43 @@ import math
 import random
 import requests
 import os
+from huggingface_hub import hf_hub_download
 
-# Dictionary to map brands to their respective pickle files
+# Dictionary to map brands to their respective HuggingFace model repo files
 BRAND_GRAPHS = {
-    'drumeo': 'drumeo_pop_items_labels.pkl',
-    'pianote': 'pianote_pop_items_labels.pkl',
-    'singeo': 'singeo_pop_items_labels.pkl',
-    'guitareo': 'guitareo_pop_items_labels.pkl'
+    'drumeo': 'drumeo_graph.pkl',
+    'pianote': 'pianote_graph.pkl',
+    'singeo': 'singeo_graph.pkl',
+    'guitareo': 'guitareo_graph.pkl'
 }
 
-# API Authorization Token
+# HuggingFace Repository Info
+#HF_REPO = "MusoraProductDepartment/popular-path-graphs"
 AUTH_TOKEN = os.getenv('HF_TOKEN')
-
 API_URL = "https://MusoraProductDepartment-PWGenerator.hf.space/rank_items/"
 
 
 @st.cache_resource
-def load_graph(brand):
-    with open(BRAND_GRAPHS[brand], 'rb') as f:
-        return pickle.load(f)
+def load_graph_from_hf(brand):
+    """
+    Load the graph for the selected brand from HuggingFace Hub.
+    """
+    try:
+        # Download the file from HuggingFace Hub
+        HF_REPO = '{brand}-graph'
+        file_path = hf_hub_download(repo_id=HF_REPO, filename=BRAND_GRAPHS[brand], use_auth_token=AUTH_TOKEN)
+        # Load the graph
+        with open(file_path, 'rb') as f:
+            return pickle.load(f)
+    except Exception as e:
+        st.error(f"Error loading graph from HuggingFace: {e}")
+        return None
+
 
 def filter_graph(graph, node_threshold=10, edge_threshold=5):
+    """
+    Filters the graph to include only popular nodes and edges.
+    """
     popular_nodes = [
         node for node in graph.nodes
         if graph.degree(node) >= node_threshold
@@ -40,6 +56,7 @@ def filter_graph(graph, node_threshold=10, edge_threshold=5):
 
     return filtered_graph
 
+
 def get_rankings_from_api(brand, user_id, content_ids):
     """
     Call the rank_items API to fetch rankings for the given user and content IDs.
@@ -68,17 +85,10 @@ def rank_to_color(rank, max_rank):
     """
     Map a rank to a grayscale color, where dark gray indicates high relevance (low rank),
     and light gray indicates low relevance (high rank).
-
-    Parameters:
-        rank (int): The rank of the item.
-        max_rank (int): The maximum rank (for normalizing the gradient).
-
-    Returns:
-        str: Hex color code for the grayscale shade.
     """
     if rank > max_rank:  # Handle items without ranking
         return "#E8E8E8"  # Very light gray for unranked items
-    intensity = int(55 + (rank / max_rank) * 200)  # Scale intensity (darker for lower ranks)
+    intensity = int(255 - (rank / max_rank) * 200)  # Darker for lower ranks
     return f"rgb({intensity}, {intensity}, {intensity})"  # Grayscale
 
 
@@ -178,14 +188,14 @@ selected_brand = st.selectbox("Select a brand:", options=list(BRAND_GRAPHS.keys(
 
 if "selected_brand" not in st.session_state or st.session_state.selected_brand != selected_brand:
     st.session_state.selected_brand = selected_brand
-    G = load_graph(selected_brand)
+    G = load_graph_from_hf(selected_brand)
 
     # Sort nodes by popularity (in-degree + out-degree) and select from top 20
     popular_nodes = sorted(G.nodes, key=lambda n: G.in_degree(n) + G.out_degree(n), reverse=True)
     top_20_nodes = popular_nodes[:20] if len(popular_nodes) > 20 else popular_nodes
     st.session_state.start_node = random.choice(top_20_nodes)
 else:
-    G = load_graph(selected_brand)
+    G = load_graph_from_hf(selected_brand)
 
 # Random Selection Button
 if st.button("Random Selection"):
@@ -201,7 +211,6 @@ start_node = st.number_input(
 # Input: Student ID
 student_id = st.text_input("Enter a student ID (optional):", value="")
 
-
 # Toggle for showing content titles
 show_titles = st.checkbox("Show content titles", value=False)
 
@@ -217,7 +226,7 @@ if student_id:
     rankings = get_rankings_from_api(selected_brand, int(student_id), content_ids)
     if rankings:
         rankings = rankings['ranked_content_ids']
-#print(rankings)
+
 layers = st.slider("Depth to explore:", 1, 6, value=3)
 top_k = st.slider("Branching factor (per node):", 1, 6, value=3)