app.py

import gradio as gr
import os
import numpy as np
import ujson as json
import time
from threading import Thread

from loading import load_data, save_git, load_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 and for security reasons cannot execute remote code. Please check the model architecture and whether it too recent and requires the execution of custom code.")
            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 or might be gated. Please check the model name and its accessibility.")
            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 max). Please verify the size of the model.")
                return None
            else:
                print(f"Error loading model '{path}': {e}")
                gr.Warning("Model couldn't be loaded. Check the logs for what happened or report an issue including the model's name.")
                return None
    except Exception as e:
            print(f"Error loading model '{path}': {e}")
            gr.Warning("Model couldn't be loaded. Check the logs for what happened or report an issue including the model's name.")
            return None
    progress_bar(0.25, desc="Generating data...",total=100)
    time0 = time.perf_counter()
    for i,output in enumerate(llm_run(model,tokenizer,genes,N)):
        time_elapsed = time.perf_counter()-time0
        estimated_time_remaining = int(len(genes)*time_elapsed/(i+1))
        minutes = str(estimated_time_remaining//60)
        minutes = "0"*(2-min(2,len(minutes))) + minutes
        seconds = str(estimated_time_remaining%60)
        seconds = "0"*(2-min(2,len(seconds))) + seconds
        progress_bar(0.25 + i*(70/len(genes))/100, desc=f"Generating data... {i+1}/{len(genes)} - estimated remaining time {minutes}:{seconds}",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 or might be gated. 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 or might be gated. 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)]
    fsave = False
    for i in range(10): #Trying to push
        try:
            save_git(alleles,genes,path,family)
            fsave = True
            break
        except Exception as e:
            print(f"Error saving data: {e}")
            #Recloning the repo
            try:
                load_git(force_clone=True)
            except Exception as e:
                print(f"Error recloning repo: {e}")
    if not fsave:
        gr.Warning("Something went wrong with GitHub and data couldn't be sent to the server. Please check the logs or report an issue. ")
        """def download_data():
            d = {'family':family,'alleles':alleles}
            model_name = path
            data_path = f'math/{model_name}.json'
            path = os.path.join('Data',data_path)
            #create the file folder path
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path), exist_ok=True)
            #Open the file
            with open(path,'w') as f:
                json.dump(d,f)
            # Provide the download link
            return gr.File.update(value=path, label="Download data", file_name=f"{model_name}.json")
        gr.Button("Download data",variant="primary").click(fn=download_data, inputs=[], outputs=None)"""
        return None
    

    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
    global ORDERED_MODEL_NAMES
    
    # 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


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("_This space is under active development. New features and improvements will be added regularly and some things may not work properly. Feel free to open an issue if you encounter any problems._")
    
    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"
        "Models are colored according to their family (e.g., LLaMA, OPT, Mistral) for the ones that were in the original paper. Models added by users are colored in grey for now. "
    )
    
    # 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("## Submit 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., `'Qwen/Qwen2.5-7B-Instruct'`).\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.\n"
    )
    
    with gr.Group():
        model_input = gr.Textbox(label="Model", interactive=True)
        submit_btn = gr.Button("Run PhyloLM", variant="primary",interactive=True)

    
    # 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)
    gr.Markdown('Socials : [BlueSky](https://bsky.app/profile/nicolasyax.bsky.social) [Twitter](https://x.com/nicolas__yax?s=11)')
    
    # 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
    run_event = 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])
    #cancel_btn.click(fn=None,inputs=None,outputs=None,cancels=[run_event])

    #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()