Update app.py

app.py

@@ -4,6 +4,8 @@ import requests
 import pandas as pd
 from io import StringIO
 import plotly.graph_objs as go
+from huggingface_hub import HfApi
+from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError
 
 from yall import create_yall
 
@@ -34,7 +36,29 @@ def convert_markdown_table_to_dataframe(md_content):
 
     return df
 
+@st.cache_data 
+def get_model_info(df):
+    api = HfApi()
+    
+    # Initialize new columns for likes and tags
+    df['likes'] = None
+    df['tags'] = None
+
+    # Iterate through DataFrame rows
+    for index, row in df.iterrows():
+        model = row['Model'].strip()
+        try:
+            model_info = api.model_info(repo_id=str(model))
+            df.loc[index, 'Likes'] = model_info.likes
+            df.loc[index, 'Tags'] = ', '.join(model_info.tags)
+
+        except (RepositoryNotFoundError, RevisionNotFoundError):
+            df.loc[index, 'Likes'] = -1
+            df.loc[index, 'Tags'] = ''
+
+    return df
 
+    
 def create_bar_chart(df, category):
     """Create and display a bar chart for a given category."""
     st.write(f"### {category} Scores")
@@ -74,25 +98,35 @@ def main():
                 score_columns = ['Average', 'AGIEval', 'GPT4All', 'TruthfulQA', 'Bigbench']
     
                 # Display dataframe
-                df = convert_markdown_table_to_dataframe(content)
+                full_df = convert_markdown_table_to_dataframe(content)
+                full_df = get_model_info(full_df)
                 for col in score_columns:
-                    df[col] = pd.to_numeric(df[col].str.strip(), errors='coerce')
-
-                # Toggles for Phi and Mistral in a single row
-                col1, col2 = st.columns(2)
+                    full_df[col] = full_df.to_numeric(full_df[col].str.strip(), errors='coerce')
+                df = pd.DataFrame(columns=full_df.columns)
+                
+                # Toggles
+                col1, col2, col3 = st.columns(3)
                 with col1:
-                    show_phi = st.checkbox("Phi", value=True)
+                    show_phi = st.toggle("Phi (2.8B)")
                 with col2:
-                    show_mistral = st.checkbox("Mistral", value=True)
+                    show_mistral = st.toggle("Mistral (7B)")
+                with col3:
+                    show_mistral = st.toggle("Mixtral (46.7B)")
 
                 # Apply filters based on toggles
-                if not show_phi:
-                    df = df[~df['Model'].str.lower().str.contains('phi')]
-                if not show_mistral:
-                    df = df[~df['Model'].str.lower().str.contains('mistral')]
+                if show_phi:
+                    df = df.append(full_df[full_df['tags'].str.lower().str.contains('phi-msft')], ignore_index=True)
+                if show_mistral:
+                    df = df.append(full_df[full_df['tags'].str.lower().str.contains('mistral')], ignore_index=True)
+                if show_mixtral:
+                    df = df.append(full_df[full_df['tags'].str.lower().str.contains('mixtral')], ignore_index=True)
+
+                # Sort values
+                df = df.sort_values(by='Average', ascending=False)
+
+                # Display the DataFrame
+                st.dataframe(df[['Model'] + score_columns + ['Likes']], use_container_width=True)
 
-                st.dataframe(df, use_container_width=True)
-    
                 # Full-width plot for the first category
                 create_bar_chart(df, score_columns[0])