app.py

import streamlit as st
import pandas as pd
from huggingface_hub import HfApi
import requests
import re
import concurrent.futures
import logging
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

logging.basicConfig(level=logging.WARNING, format="%(asctime)s - %(levelname)s - %(message)s")

# Configure the Streamlit page
st.set_page_config(page_title="Heretic Models Explorer", page_icon="🔥", layout="wide")

st.title("🔥 Heretic Models Explorer")
st.markdown(
    "This space lists all models on Hugging Face tagged with "
    "[`heretic`](https://huggingface.co/models?other=heretic). "
    "It automatically fetches their model cards to extract **KL Divergence** and **Refusals**, "
    "allowing you to sort and compare them easily. Click on any column header to sort!"
)

@retry(
    retry=retry_if_exception_type((requests.exceptions.ConnectionError, requests.exceptions.Timeout, requests.exceptions.HTTPError)),
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=2, max=10)
)
def fetch_url(url):
    response = requests.get(url, timeout=5)
    # Raise HTTPError for rate limits (429) or server errors (5xx) to trigger a retry
    if response.status_code in (429, 500, 502, 503, 504):
        response.raise_for_status()
    return response

# Fetch model's readme.md file and attempt to parse refusal and KL divergence data.
def fetch_model_info(model):
    """Fetches the README.md for a given model and dynamically extracts metrics."""
    model_id = model.id
    kl_div = None
    refusals_str = None
    refusal_rate = None
    initial_refusals_str = None
    
    # Download the README.md via raw URL for high-speed fetching
    url = f"https://huggingface.co/{model_id}/raw/main/README.md"
    try:
        response = fetch_url(url)
        if response.status_code == 200:
            readme_text = response.text
            
            # Extract KL divergence using Regex
            # Matches formats like: "KL divergence | 0.0033"
            kl_match = re.search(r"(?i)KL\s*divergence[^a-zA-Z\d\n]*?([\d\.]+)", readme_text)
            if kl_match:
                try:
                    kl_div = float(kl_match.group(1))
                except ValueError as e:
                    logging.warning(f"Could not parse KL divergence for {model_id}: {e}")
                    
            # Extract Refusals using Regex
            # Matches formats like: "Refusals | 15/100" or "15 / 100"
            for line in readme_text.split('\n'):
                line_lower = line.lower()
                if 'refusals' in line_lower:
                    fractions = re.findall(r"(\d+)\s*/\s*(\d+)", line)
                    if not fractions:
                        continue
                    
                    if 'initial' in line_lower and initial_refusals_str is None:
                        initial_refusals_str = f"{fractions[0][0]}/{fractions[0][1]}"
                    elif 'initial' not in line_lower and refusals_str is None:
                        refusals_str = f"{fractions[0][0]}/{fractions[0][1]}"
                        try:
                            refusal_rate = (int(fractions[0][0]) / int(fractions[0][1])) * 100 
                        except (ValueError, ZeroDivisionError) as e:
                            logging.warning(f"Could not parse refusals for {model_id}: {e}")
                            continue
                        if len(fractions) >= 2 and initial_refusals_str is None:
                            initial_refusals_str = f"{fractions[1][0]}/{fractions[1][1]}"
                            
    except requests.exceptions.RequestException as e:
        logging.warning(f"Error fetching README for {model_id}: {e}")
    except Exception as e:
        logging.error(f"Unexpected error parsing README for {model_id}: {e}")

    # Skip models without explicitly stated KL divergence or refusal data
    if kl_div is None or refusals_str is None:
        return None

    return {
        "Model ID": model_id,
        "KL Divergence": kl_div,
        "Initial Refusals": initial_refusals_str,
        "Refusal Rate (%)": refusal_rate,
        "Refusals": refusals_str,
        "Likes": getattr(model, 'likes', 0),
        "Downloads": getattr(model, 'downloads', 0),
        "URL": f"https://huggingface.co/{model_id}"
    }

# Check for reproducibility record first and then fallback to parsing data from readme.md 
def get_model_data(model):
    model_id = model.id
    url = f"https://huggingface.co/{model_id}/raw/main/reproduce/reproduce.json"
    try:
        response = fetch_url(url)
        if response.status_code == 200:
            data = response.json()
            metrics = data.get("metrics")
            if isinstance(metrics, dict):
                kl_div = metrics.get("kl_divergence") or metrics.get("kl")
                refusals = metrics.get("refusals")
                base_refusals = metrics.get("base_refusals")
                n_bad = metrics.get("n_bad_prompts")
                if kl_div is not None and refusals is not None and n_bad is not None:
                    initial_refusals_str = f"{base_refusals}/{n_bad}" if base_refusals is not None else None
                    refusals_str = f"{refusals}/{n_bad}"
                    refusal_rate = (refusals / n_bad) * 100 if n_bad > 0 else 0
                    return {
                        "Model ID": model_id,
                        "KL Divergence": float(kl_div),
                        "Initial Refusals": initial_refusals_str,
                        "Refusal Rate (%)": refusal_rate,
                        "Refusals": refusals_str,
                        "Likes": getattr(model, 'likes', 0),
                        "Downloads": getattr(model, 'downloads', 0),
                        "URL": f"https://huggingface.co/{model_id}"
                    }
    except requests.exceptions.RequestException as e:
        logging.warning(f"Error fetching JSON for {model_id}: {e}")
    except (ValueError, KeyError, TypeError) as e:
        logging.warning(f"Formatting error in JSON for {model_id}: {e}")
    except Exception as e:
        logging.error(f"Unexpected error processing JSON for {model_id}: {e}")
    return fetch_model_info(model)

@st.cache_data(ttl=3600, show_spinner=False) # Cache for 1 hour
def get_heretic_models():
    """Fetches all heretic models and their metrics concurrently."""
    api = HfApi()
    
    # Query all models using the Hugging Face Hub `filter` parameter
    models = list(api.list_models(filter="heretic"))
    
    # Fetch details concurrently for speed (10 workers to avoid API rate limiting issues)
    # Skip quantised models
    quant_patterns = re.compile(r'(?i)(gguf|mlx|awq|nvfp4|gptq|exl[23]|-quant|int8|int4|oq[1-8]|mxfp[48])')
    models = [m for m in models if not quant_patterns.search(m.id)]
    
    data = []
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
        results = executor.map(get_model_data, models)
        for res in results:
            if res is not None:
                data.append(res)
            
    return data

# Main execution
with st.spinner("Fetching heretic models and parsing model cards... This might take a moment on the first run."):
    models_data = get_heretic_models()

if not models_data:
    st.warning("No models found with the 'heretic' tag.")
else:
    df = pd.DataFrame(models_data)
    
    # Make Model ID a clickable Markdown link
    df["Model"] = df["URL"]
    
    # Select and order columns for display
    display_df = df[["Model", "KL Divergence", "Initial Refusals", "Refusals", "Refusal Rate (%)", "Likes", "Downloads"]]
    
    st.markdown(f"**Found {len(display_df)} models.**")
    
    # Display as an interactive, sortable dataframe
    st.dataframe(
        display_df,
        column_config={
            "Model": st.column_config.LinkColumn("Model", display_text=r"https://huggingface\.co/(.*)"),
            "KL Divergence": st.column_config.NumberColumn("KL Divergence", format="%.4f"),
            "Initial Refusals": st.column_config.TextColumn("Initial Refusals"),
            "Refusals": st.column_config.TextColumn("Refusals"),
            "Refusal Rate (%)": st.column_config.NumberColumn("Refusal Rate (%)", format="%.2f%%"),
            "Likes": st.column_config.NumberColumn("Likes"),
            "Downloads": st.column_config.NumberColumn("Downloads")
        },
        hide_index=True,
        width='stretch'
    )