app.py

import gradio as gr
import os
import numpy as np
import ujson as json

from loading import load_data, save_git
from tools import compute_ordered_matrix
from plotting import plot_sim_matrix_fig, plot_umap_fig, plot_tree, update_sim_matrix_fig, update_umap_fig, update_tree_fig
from llm_run import download_llm_to_cache, load_model, llm_run

def reload_figures():
    global MODEL_SEARCHED_X, MODEL_SEARCHED_Y, ALPHA_EDGES, ALPHA_NAMES, ALPHA_MARKERS, FIGS, ORDERED_MODEL_NAMES
    fig1 = update_sim_matrix_fig(FIGS['fig1'], ORDERED_MODEL_NAMES, model_search_x=MODEL_SEARCHED_X, model_search_y=MODEL_SEARCHED_Y)
    fig2 = update_umap_fig(FIGS['fig2'], DIST_MATRIX, MODEL_NAMES, FAMILIES, COLORS, model_search_x=MODEL_SEARCHED_X, alpha_edges=ALPHA_EDGES['fig2'], alpha_names=ALPHA_NAMES['fig2'], alpha_markers=ALPHA_MARKERS['fig2'])
    fig4 = update_tree_fig(FIGS['fig4'], MODEL_NAMES, model_search=MODEL_SEARCHED_X, alpha_edges=ALPHA_EDGES['fig4'], alpha_names=ALPHA_NAMES['fig4'], alpha_markers=ALPHA_MARKERS['fig4'])
    return [fig1,fig2,fig4]

def search_bar_changeX(value):
    global MODEL_SEARCHED_X
    MODEL_SEARCHED_X = value
    return reload_figures()

def search_bar_changeY(value):
    global MODEL_SEARCHED_Y
    MODEL_SEARCHED_Y = value
    return reload_figures()

def slider_changeAlphaMarkers(value,key):
    global ALPHA_MARKERS
    ALPHA_MARKERS[key] = value
    return reload_figures()

def slider_changeAlphaNames(value,key):
    global ALPHA_NAMES
    ALPHA_NAMES[key] = value
    return reload_figures()

def slider_changeAlphaEdges(value,key):
    global ALPHA_EDGES
    ALPHA_EDGES[key] = value
    return reload_figures()

def search_bar_gr(model_names,slider=True,double_search=False,key=None):
    global MODEL_SEARCHED_X,MODEL_SEARCHED_Y,ALPHA_EDGES,ALPHA_NAMES, ALPHA_MARKERS
    #col1,col2 = gr.Row([0.2,0.8])
    ret = []
    with gr.Column(scale=1) as col1:
        with gr.Group():
            if MODEL_SEARCHED_X is None:
                index = 0
            else:
                index = model_names.index(MODEL_SEARCHED_X)
            ms_x = gr.Dropdown(label='Search'+(' X' if double_search else ''),choices=model_names,value=model_names[index],key='model_search_x_'+key,interactive=True)
            #set MODEL_SEARCH_X
            ret.append(ms_x)
            if double_search:
                if MODEL_SEARCHED_Y is None:
                    index = 0
                else:
                    index = model_names.index(MODEL_SEARCHED_Y)
                ms_y = gr.Dropdown(label='Search Y',choices=model_names,value=model_names[index],key='model_search_y_'+key,interactive=True)
                ret.append(ms_y)
        if slider:
            with gr.Group():
                values = np.arange(0, 1.05,0.05)
                #truncate values to the 100th
                values = np.round(values,2)
                alpha_edges = gr.Slider(label='Alpha Edges',
                                minimum=0,
                                maximum=1,
                                step=0.05,
                                value=ALPHA_EDGES[key],
                                    key='alpha_edges_'+key,
                                    interactive=True)
                
                values = np.arange(0, 1.05,0.05)
                #truncate values to the 100th
                values = np.round(values,2)
                alpha_names = gr.Slider(label='Alpha Names',
                                minimum=0,
                                maximum=1,
                                step=0.05,
                                value=ALPHA_NAMES[key],
                                    key='alpha_names_'+key,
                                    interactive=True)
                
                values = np.arange(0, 1.05,0.05)
                #truncate values to the 100th
                values = np.round(values,2)
                alpha_markers = gr.Slider(label='Alpha Markers',
                                minimum=0,
                                maximum=1,
                                step=0.05,
                                value=ALPHA_MARKERS[key],
                                    key='alpha_markers_'+key,
                                    interactive=True)
                ret.append(alpha_edges)
                ret.append(alpha_names)
                ret.append(alpha_markers)
    col2 = gr.Column(scale=5)
    ret.insert(0,col2)
    return ret

import spaces
@spaces.GPU(duration=300)
def _run(path,genes,N,progress_bar):
    #Load the model
    progress_bar(0.20, desc="Loading Model...",total=100)
    try:
        model,tokenizer = load_model(path)
    except ValueError as e:
            print(f"Error loading model '{path}': {e}")
            gr.Warning("Model couldn't load. This space currently only works with AutoModelForCausalLM models. Please check the model architecture and try again.")
            return None
    except OSError as e:
            print(f"Error loading model '{path}': {e}")
            gr.Warning("Model doesn't seem to exist on the HuggingFace Hub. Please check the model name and try again.")
            return None
    except RuntimeError as e:
            if 'out of memory' in str(e):
                print(f"Error loading model '{path}': {e}")
                gr.Warning("Loading the model triggered an out of memory error. It may be too big for the GPU (80Go RAM). Please try again with a smaller model.")
                return None
            else:
                print(f"Error loading model '{path}': {e}")
                gr.Warning("Model couldn't be loaded. Please check the logs or report an issue.")
                return None
    except Exception as e:
            print(f"Error loading model '{path}': {e}")
            gr.Warning("Model couldn't be loaded. Please check logs or report an issue.")
            return None
    progress_bar(0.25, desc="Generating data...",total=100)
    for i,output in enumerate(llm_run(model,tokenizer,genes,N)):
        progress_bar(0.25 + i*(70/len(genes))/100, desc=f"Generating data... {i+1}/{len(genes)}",total=100)
    return output

def run(path,progress_bar):
    global DEFAULT_FAMILY_NAME, PHYLOLM_N
    family = DEFAULT_FAMILY_NAME
    N = PHYLOLM_N
    #Loading bar
    progress_bar(0, desc="Downloading model...",total=100)
    try:
        # Download the model to cache
        if download_llm_to_cache(path) is None:
            gr.Warning("Model not found on Hugging Face Hub. Please check the model name and try again.")
            return None
    except OSError as e:
        print(f"Error downloading model: {e}")
        gr.Warning("Model not found on Hugging Face Hub. Please check the model name and try again.")
        return None
    
    # Load the model
    progress_bar(0.10, desc="Loading contexts...",total=100)

    with open('inputs/math.json', 'r') as f:
        genes = json.load(f)

    # Load the model and run
    progress_bar(0.15, desc="Waiting for GPU...",total=100)

    try:
        output = _run(path,genes,N,progress_bar)
        if output is None:
            return None
    except Exception as e:
        print(f"Error running model: {e}")
        gr.Warning("Something unexpected happened during the run or the loading of the model. Please check the logs or report an issue.")
        return None

    progress_bar(0.95, desc="Saving data ...",total=100)

    alleles = [[compl[j]['generated_text'][len(gene):][:4] for j in range(len(compl))] for gene,compl in zip(genes,output)]
    save_git(alleles,genes,path,family)

    progress_bar(1, desc="Done!",total=100)


def prepare_run(model_name,progress_bar=gr.Progress()):
    global MODEL_SEARCHED_X,MODEL_NAMES
    if model_name in MODEL_NAMES:
        gr.Warning('Model already exists in the database.')
        MODEL_SEARCHED_X = model_name
        reload_figures()
        return
    run(model_name,progress_bar)

def reload_env():
    global SIM_MAT_SEARCH_X, SIM_MAT_SEARCH_Y, VIZ_SEARCH, TREE_SEARCH
    global MODEL_NAMES, FAMILIES, COLORS, SIM_MATRIX, DIST_MATRIX
    global FIGS, FIGS_OBJECTS
    
    # Load models for the dropdown
    data, model_names, families, sim_matrix, colors = load_data()

    sim_matrix_safe = np.where(sim_matrix == 0, np.finfo(np.float64).eps, sim_matrix)
    dist_matrix = -np.log(sim_matrix_safe)

    #Set globals
    MODEL_NAMES = model_names
    FAMILIES = families
    COLORS = colors
    SIM_MATRIX = sim_matrix
    DIST_MATRIX = dist_matrix

    #Update Figs
    ordered_sim_matrix, ordered_model_names = compute_ordered_matrix(sim_matrix,dist_matrix, model_names)
    ORDERED_MODEL_NAMES = ordered_model_names
    FIGS['fig1'] = plot_sim_matrix_fig(ordered_sim_matrix, ordered_model_names, families, colors)
    FIGS['fig2'] = plot_umap_fig(dist_matrix, sim_matrix, model_names, families, colors,
                                        alpha_edges=ALPHA_EDGES['fig2'],alpha_names=ALPHA_NAMES['fig2'],alpha_markers=ALPHA_MARKERS['fig2'])
    FIGS['fig4'] = plot_tree(sim_matrix, model_names, families, colors,alpha_edges=ALPHA_EDGES['fig4'],alpha_names=ALPHA_NAMES['fig4'],alpha_markers=ALPHA_MARKERS['fig4'])

    #Update search bars
    sim_mat_search_x = gr.Dropdown(label='Search X',choices=model_names,value=model_names[0],key='model_search_x_fig1',interactive=True)
    sim_mat_search_y = gr.Dropdown(label='Search Y',choices=model_names,value=model_names[0],key='model_search_y_fig1',interactive=True)
    viz_search = gr.Dropdown(label='Search',choices=model_names,value=model_names[0],key='model_search_fig2',interactive=True)
    tree_search = gr.Dropdown(label='Search',choices=model_names,value=model_names[0],key='model_search_fig4',interactive=True)

    return FIGS['fig1'], FIGS['fig2'], FIGS['fig4'], sim_mat_search_x, sim_mat_search_y, viz_search, tree_search

    

# Load environment variables

USERNAME = os.environ['GITHUB_USERNAME']
TOKEN = os.environ['GITHUB_TOKEN']
MAIL = os.environ['GITHUB_MAIL']

MODEL_SEARCHED_X = None
MODEL_SEARCHED_Y = None
ALPHA_EDGES = {'fig2':0.05, 'fig3':0.05,'fig4':1.0}
ALPHA_NAMES = {'fig2':0.0, 'fig3':0.0,'fig4':0.0}
ALPHA_MARKERS = {'fig2':0.8, 'fig3':0.8,'fig4':1.0}

FIGS = {'fig1':None,'fig2':None,'fig3':None,'fig4':None}
FIGS_OBJECTS = [None,None,None]
MODEL_NAMES = None
FAMILIES = None
COLORS = None
ORDERED_MODEL_NAMES = None
SIM_MATRIX = None
DIST_MATRIX = None

DEFAULT_FAMILY_NAME = '?'
PHYLOLM_N = 32

SIM_MAT_SEARCH_X = None
SIM_MAT_SEARCH_Y = None
VIZ_SEARCH = None
TREE_SEARCH = None

# Build the Gradio interface
with gr.Blocks(title="PhyloLM", theme=gr.themes.Default()) as demo:
    gr.Markdown("# PhyloLM: Phylogenetic Mapping of Language Models")
    
    gr.Markdown(
        "Welcome to PhyloLM ([paper](https://arxiv.org/abs/2404.04671) - [code](https://github.com/Nicolas-Yax/PhyloLM)) — a tool for comparing language models based on their **behavioral similarity**, inspired by methods from comparative genomics. "
        "Instead of architecture or weights, we use output behavior on diagnostic prompts as a behavioral fingerprint to compute a distance metric, akin to how biologists compare species using genetic data. This makes it possible to draw a unique map of all LLMs (various architectures, gated and non gated, ...)."
        "The goal of this space is to create a collaborative space where everyone can visualize these maps and extend them with models of their choice. "
    )
    
    gr.Markdown("## Explore Maps of Models")
    
    gr.Markdown(
        "This interactive space allows users to explore model similarities through four types of visualizations:\n"
        "- A similarity matrix (values range from 0 = dissimilar to 1 = highly similar). \n"
        "- 2D and 3D scatter plots representing how close or far from each other LLMs are (plotted using UMAP). \n"
        "- A tree to visualize distances between models (distance from leaf A to leaf B in the tree is similar to the distance between the two models)\n\n"
    )
    
    # Load models for the dropdown
    data, model_names, families, sim_matrix, colors = load_data()

    sim_matrix_safe = np.where(sim_matrix == 0, np.finfo(np.float64).eps, sim_matrix)
    dist_matrix = -np.log(sim_matrix_safe)

    #Set globals
    MODEL_NAMES = model_names
    FAMILIES = families
    COLORS = colors
    SIM_MATRIX = sim_matrix
    DIST_MATRIX = dist_matrix

    # Create the tabs
    tab_state = gr.State(value="Similarity Matrix")  # Default tab
    tabs = gr.Tabs(["Similarity Matrix", "2D Visualization","Tree Visualization"])
    with tabs:
        with gr.TabItem("Similarity Matrix"):
            # Similarity matrix visualization
            with gr.Row():
                col2,sim_mat_search_x,sim_mat_search_y = search_bar_gr(model_names,slider=False,double_search=True,key='fig1')
                with col2:
                    ordered_sim_matrix, ordered_model_names = compute_ordered_matrix(sim_matrix,dist_matrix, model_names)
                    fig = plot_sim_matrix_fig(ordered_sim_matrix, ordered_model_names, families, colors)
                    sim_matrix_output = gr.Plot(fig,label="Similarity Matrix")
                    FIGS['fig1'] = fig
                    ORDERED_MODEL_NAMES = ordered_model_names
                    FIGS_OBJECTS[0] = sim_matrix_output
        with gr.TabItem("2D Visualization"):
            # 2D visualization
            with gr.Row():
                col2,viz_search,viz_alpha_edge,viz_alpha_name,viz_alpha_marker = search_bar_gr(model_names,slider=True,double_search=False,key='fig2')
                with col2:
                    fig = plot_umap_fig(dist_matrix, sim_matrix, model_names, families, colors,
                                        alpha_edges=ALPHA_EDGES['fig2'],alpha_names=ALPHA_NAMES['fig2'],alpha_markers=ALPHA_MARKERS['fig2'])
                    plot_output = gr.Plot(fig,label="2D Visualization")
                    FIGS['fig2'] = fig
                    FIGS_OBJECTS[1] = plot_output
        with gr.TabItem("Tree Visualization"):
            # Tree visualization
            with gr.Row():
                col2,tree_search,tree_alpha_edge,tree_alpha_name,tree_alpha_marker = search_bar_gr(model_names,slider=True,double_search=False,key='fig4')
                with col2:
                    fig = plot_tree(sim_matrix, model_names, families, colors,alpha_edges=ALPHA_EDGES['fig4'],alpha_names=ALPHA_NAMES['fig4'],alpha_markers=ALPHA_MARKERS['fig4'])
                    tree_output = gr.Plot(fig,label="Tree Visualization")
                    FIGS['fig4'] = fig
                    FIGS_OBJECTS[2] = tree_output
            

    # Submit model section
    gr.Markdown("## Submitting a Model")
    
    gr.Markdown(
        "You may contribute new models to this collaborative space using compute resources. "
        "Once processed, the model will be compared to existing ones, and its results added to a shared public database. "
        "Model families (e.g., LLaMA, OPT, Mistral) are extracted from Hugging Face model cards and used only for visualization (e.g., coloring plots); they are **not** involved in the computation of similarity."
    )
    
    gr.Markdown(
        "**To add a new model:**\n"
        "1. Enter the name of a model hosted on Hugging Face (e.g., `'mistralai/Mistral-7B-Instruct-v0.3'`).\n"
        "2. Click on the **Run PhyloLM** button.\n"
        "- If the model has already been processed, you'll be notified and no new run will start.\n"
        "- If it hasn't been processed, it will be downloaded and be evaluated.\n\n"
        "⚠️ Be careful when submitting large LLMs (typically >15B parameters) as they may exceed the GPU RAM or the time limit, leading to failed runs."
    )
    
    with gr.Group():
        model_input = gr.Textbox(label="Model", interactive=True)
        submit_btn = gr.Button("Run PhyloLM", variant="primary")

    
    # Disclaimer and citation
    gr.Markdown("## Disclaimer")
    gr.Markdown(
        "This is a research prototype and may contain bugs or limitations. "
        "All computed data are public and hosted on [GitHub](https://github.com/PhyloLM/Data). "
        "If you'd like to contribute additional models — especially for gated or large models that cannot be processed via the web interface — "
        "you are welcome to submit a pull request to the repository cited above. "
        "All results are computed on the 'Math' set of genes used in the original paper."
    )
    
    gr.Markdown("## Citation")
    gr.Markdown("If you find this project useful for your research, please consider citing the following paper:")

    #bibtex
    gr.Code('''@inproceedings{
yax2025phylolm,
title={Phylo{LM}: Inferring the Phylogeny of Large Language Models and Predicting their Performances in Benchmarks},
author={Nicolas Yax and Pierre-Yves Oudeyer and Stefano Palminteri},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=rTQNGQxm4K}
}''',language=None)
    
    # Change actions from search bars
    sim_mat_search_x.change(fn=search_bar_changeX, inputs=sim_mat_search_x, outputs=FIGS_OBJECTS)
    sim_mat_search_y.change(fn=search_bar_changeY, inputs=sim_mat_search_y, outputs=FIGS_OBJECTS)

    viz_search.change(fn=search_bar_changeX, inputs=viz_search, outputs=FIGS_OBJECTS)

    tree_search.change(fn=search_bar_changeX, inputs=tree_search, outputs=FIGS_OBJECTS)

    # Change actions from sliders
    viz_alpha_edge.change(fn=lambda x : slider_changeAlphaEdges(x,'fig2'), inputs=viz_alpha_edge, outputs=FIGS_OBJECTS)
    viz_alpha_name.change(fn=lambda x : slider_changeAlphaNames(x,'fig2'), inputs=viz_alpha_name, outputs=FIGS_OBJECTS)
    viz_alpha_marker.change(fn=lambda x : slider_changeAlphaMarkers(x,'fig2'), inputs=viz_alpha_marker, outputs=FIGS_OBJECTS)

    tree_alpha_edge.change(fn=lambda x : slider_changeAlphaEdges(x,'fig4'), inputs=tree_alpha_edge, outputs=FIGS_OBJECTS)
    tree_alpha_name.change(fn=lambda x : slider_changeAlphaNames(x,'fig4'), inputs=tree_alpha_name, outputs=FIGS_OBJECTS)
    tree_alpha_marker.change(fn=lambda x : slider_changeAlphaMarkers(x,'fig4'), inputs=tree_alpha_marker, outputs=FIGS_OBJECTS)

    # Run PhyloLM button
    submit_btn.click(fn=prepare_run, inputs=[model_input], outputs=[model_input]).then(fn=reload_env, inputs=[], outputs=FIGS_OBJECTS+ [sim_mat_search_x, sim_mat_search_y, viz_search, tree_search])

    #Set more globals
    SIM_MAT_SEARCH_X = sim_mat_search_x
    SIM_MAT_SEARCH_Y = sim_mat_search_y
    VIZ_SEARCH = viz_search
    TREE_SEARCH = tree_search

    
    
if __name__ == "__main__":
    demo.launch()