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 and trust_remote_code=False. Please check the model architecture and whether it requires the execution of custom code 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 or might be gated. 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 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)]
    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
from dotenv import load_dotenv
load_dotenv()

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"
    )
    
    # 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., `'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."
    )
    
    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()