File size: 6,294 Bytes
26cdd43
 
 
 
 
 
9565da9
26cdd43
 
 
fab1822
 
26cdd43
171cab8
 
4b096db
171cab8
4b096db
 
171cab8
 
26cdd43
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fab1822
26cdd43
fab1822
26cdd43
fab1822
 
 
 
26cdd43
fab1822
26cdd43
 
171cab8
26cdd43
 
 
bedfad8
26cdd43
 
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
import gradio as gr

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>.</p>
"""


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


def get_license_color(model):
    """Get the color of the model based on its license."""
    try:
        card = ModelCard.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'


def get_model_names(model, genealogy, found_models=None):
    """Get a list of parent model names from the model id."""
    model_tags = []

    if found_models is None:
        found_models = []

    try:
        card = ModelCard.load(model)
        card_dict = card.data.to_dict()  # Convert the ModelCard object to a dictionary
        license = card_dict['license']

        # Check the base_model in metadata
        if 'base_model' in card_dict:
            model_tags = card_dict['base_model']

        # Check the tags in metadata
        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]

        # Check for merge.yml and mergekit_config.yml if no model_tags found in the tags
        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"))

        # Convert to a list if tags is not None or empty, else set to an empty list
        if not isinstance(model_tags, list):
            model_tags = [model_tags] if model_tags else []

        # Add found model names to the list
        found_models.extend(model_tags)

        # Record the genealogy
        for model_tag in model_tags:
            genealogy[model_tag].append(model)

        # Recursively check for more models
        for model_tag in model_tags:
            get_model_names(model_tag, genealogy, found_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


def create_family_tree(start_model):
    genealogy = defaultdict(list)
    get_model_names(start_model, genealogy)  # Assuming this populates the genealogy

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

    # Get max depth
    max_depth = nx.dag_longest_path_length(G) + 1

    # Get max width
    max_width = max_width_of_tree(G) + 1

    # 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/NeuralBeagle-7B")
    btn = gr.Button("Create tree")
    out = gr.Image()
    btn.click(fn=create_family_tree, inputs=model_id, outputs=out)

demo.launch()