File size: 7,266 Bytes
26cdd43
47913ac
26cdd43
 
 
 
9565da9
26cdd43
 
 
fab1822
 
5f8cec7
 
 
 
 
 
 
 
490def8
 
26cdd43
490def8
47913ac
548eda9
 
 
 
 
 
 
 
 
 
 
 
 
47913ac
490def8
26cdd43
5f8cec7
26cdd43
 
 
 
 
 
5f8cec7
490def8
26cdd43
5f8cec7
26cdd43
47913ac
26cdd43
5f8cec7
 
 
26cdd43
5f8cec7
26cdd43
5f8cec7
26cdd43
5f8cec7
26cdd43
 
548eda9
490def8
47913ac
5f8cec7
 
26cdd43
47913ac
 
 
 
 
5f8cec7
47913ac
 
26cdd43
 
47913ac
 
26cdd43
 
47913ac
26cdd43
 
 
 
 
 
 
 
 
 
 
 
 
 
 
47913ac
26cdd43
 
 
47913ac
26cdd43
 
5f8cec7
26cdd43
 
 
5f8cec7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
548eda9
490def8
26cdd43
 
5f8cec7
 
 
26cdd43
5f8cec7
26cdd43
 
5f8cec7
26cdd43
 
 
 
5f8cec7
 
 
 
 
 
 
 
 
 
26cdd43
 
 
5f8cec7
26cdd43
 
 
5f8cec7
26cdd43
 
5f8cec7
26cdd43
 
5f8cec7
26cdd43
 
5f8cec7
26cdd43
 
 
 
 
 
 
 
548eda9
5f8cec7
 
 
 
 
 
 
 
 
 
548eda9
5f8cec7
 
 
 
 
 
 
 
 
 
548eda9
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
import gradio as gr
import sys
from huggingface_hub import ModelCard, HfApi
import requests
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from collections import defaultdict
from networkx.drawing.nx_pydot import graphviz_layout
from io import BytesIO
from PIL import Image

TITLE = """
<div align="center">
    <p style="font-size: 36px;">🌳 Model Family Tree</p>
</div><br/>
<p>Automatically calculate the <strong>family tree of a given model</strong>. It also displays the type of license each model uses (permissive, noncommercial, or unknown).</p>
<p>You can also run the code in this <a href="https://colab.research.google.com/drive/1s2eQlolcI1VGgDhqWIANfkfKvcKrMyNr?usp=sharing">Colab notebook</a>. Special thanks to <a href="https://huggingface.co/leonardlin">leonardlin</a> for his caching implementation. See also mrfakename's version in <a href="https://huggingface.co/spaces/mrfakename/merge-model-tree">this space</a>.</p>
"""

# Define the model ID
MODEL_ID = "mlabonne/NeuralBeagle14-7B"

# Define a class to cache model cards
class CachedModelCard(ModelCard):
    _cache = {}

    @classmethod
    def load(cls, model_id: str, **kwargs) -> "ModelCard":
        if model_id not in cls._cache:
            try:
                print('REQUEST ModelCard:', model_id)
                cls._cache[model_id] = super().load(model_id, **kwargs)
            except:
                cls._cache[model_id] = None
        else:
            print('CACHED:', model_id)
        return cls._cache[model_id]

# Function to get model names from a YAML file
def get_model_names_from_yaml(url):
    """Get a list of parent model names from the yaml file."""
    model_tags = []
    response = requests.get(url)
    if response.status_code == 200:
        model_tags.extend([item for item in response.content if '/' in str(item)])
    return model_tags


# Function to get the color of the model based on its license
def get_license_color(model):
    """Get the color of the model based on its license."""
    try:
        card = CachedModelCard.load(model)
        license = card.data.to_dict()['license'].lower()
        # Define permissive licenses
        permissive_licenses = ['mit', 'bsd', 'apache-2.0', 'openrail']  # Add more as needed
        # Check license type
        if any(perm_license in license for perm_license in permissive_licenses):
            return 'lightgreen'  # Permissive licenses
        else:
            return 'lightcoral'  # Noncommercial or other licenses
    except Exception as e:
        print(f"Error retrieving license for {model}: {e}")
        return 'lightgray'


# Function to find model names in the family tree
def get_model_names(model, genealogy, found_models=None, visited_models=None):
    print('---')
    print(model)
    if found_models is None:
        found_models = set()
    if visited_models is None:
        visited_models = set()

    if model in visited_models:
        print("Model already visited...")
        return found_models
    visited_models.add(model)

    try:
        card = CachedModelCard.load(model)
        card_dict = card.data.to_dict()
        license = card_dict['license']

        model_tags = []
        if 'base_model' in card_dict:
            model_tags = card_dict['base_model']

        if 'tags' in card_dict and not model_tags:
            tags = card_dict['tags']
            model_tags = [model_name for model_name in tags if '/' in model_name]

        if not model_tags:
            model_tags.extend(get_model_names_from_yaml(f"https://huggingface.co/{model}/blob/main/merge.yml"))
        if not model_tags:
            model_tags.extend(get_model_names_from_yaml(f"https://huggingface.co/{model}/blob/main/mergekit_config.yml"))

        if not isinstance(model_tags, list):
            model_tags = [model_tags] if model_tags else []

        found_models.add(model)

        for model_tag in model_tags:
            genealogy[model_tag].append(model)
            get_model_names(model_tag, genealogy, found_models, visited_models)

    except Exception as e:
        print(f"Could not find model names for {model}: {e}")

    return found_models

def find_root_nodes(G):
    """ Find all nodes in the graph with no predecessors """
    return [n for n, d in G.in_degree() if d == 0]


def max_width_of_tree(G):
    """ Calculate the maximum width of the tree """
    max_width = 0
    for root in find_root_nodes(G):
        width_at_depth = calculate_width_at_depth(G, root)
        local_max_width = max(width_at_depth.values())
        max_width = max(max_width, local_max_width)
    return max_width


def calculate_width_at_depth(G, root):
    """ Calculate width at each depth starting from a given root """
    depth_count = defaultdict(int)
    queue = [(root, 0)]
    while queue:
        node, depth = queue.pop(0)
        depth_count[depth] += 1
        for child in G.successors(node):
            queue.append((child, depth + 1))
    return depth_count


# Function to create the family tree
def create_family_tree(start_model):
    genealogy = defaultdict(list)
    get_model_names(start_model, genealogy)  # Assuming this populates the genealogy

    print("Number of models:", len(CachedModelCard._cache))

    # Create a directed graph
    G = nx.DiGraph()

    # Add nodes and edges to the graph
    for parent, children in genealogy.items():
        for child in children:
            G.add_edge(parent, child)

    try:
        # Get max depth and width
        max_depth = nx.dag_longest_path_length(G) + 1
        max_width = max_width_of_tree(G) + 1
    except:
        # Get max depth and width
        max_depth = 21
        max_width = 9

    # Estimate plot size
    height = max(8, 1.6 * max_depth)
    width = max(8, 6 * max_width)

    # Set Graphviz layout attributes for a bottom-up tree
    plt.figure(figsize=(width, height))
    pos = graphviz_layout(G, prog="dot")

    # Determine node colors based on license
    node_colors = [get_license_color(node) for node in G.nodes()]

    # Create a label mapping with line breaks
    labels = {node: node.replace("/", "\n") for node in G.nodes()}

    # Draw the graph
    nx.draw(G, pos, labels=labels, with_labels=True, node_color=node_colors, font_size=12, node_size=8_000, edge_color='black')

    # Create a legend for the colors
    legend_elements = [
        Patch(facecolor='lightgreen', label='Permissive'),
        Patch(facecolor='lightcoral', label='Noncommercial'),
        Patch(facecolor='lightgray', label='Unknown')
    ]
    plt.legend(handles=legend_elements, loc='upper left')

    plt.title(f"{start_model}'s Family Tree", fontsize=20)

    # Capture the plot as an image in memory
    img_buffer = BytesIO()
    plt.savefig(img_buffer, format='png', bbox_inches='tight')
    plt.close()
    img_buffer.seek(0)

    # Open the image using PIL
    img = Image.open(img_buffer)

    return img

with gr.Blocks() as demo:
    gr.Markdown(TITLE)
    model_id = gr.Textbox(label="Model ID", value="mlabonne/NeuralBeagle14-7B")
    btn = gr.Button("Create tree")
    out = gr.Image()
    btn.click(fn=create_family_tree, inputs=model_id, outputs=out)

demo.queue(api_open=False).launch(show_api=False)


create_family_tree(MODEL_ID)