Spaces:
Running
Running
| # Import necessary libraries | |
| 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 | |
| # Define the model ID | |
| MODEL_ID = "mlabonne/NeuralBeagle14-7B" | |
| # Define a class to cache model cards | |
| class CachedModelCard(ModelCard): | |
| _cache = {} | |
| def load(cls, model_id: str, **kwargs) -> "ModelCard": | |
| if model_id not in cls._cache: | |
| try: | |
| cls._cache[model_id] = super().load(model_id, **kwargs) | |
| except: | |
| cls._cache[model_id] = None | |
| return cls._cache[model_id] | |
| # Function to get model names from a YAML file | |
| def get_model_names_from_yaml(url): | |
| 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): | |
| try: | |
| card = CachedModelCard.load(model) | |
| license = card.data.to_dict()['license'].lower() | |
| permissive_licenses = ['mit', 'bsd', 'apache-2.0', 'openrail'] | |
| if any(perm_license in license for perm_license in permissive_licenses): | |
| return 'lightgreen' | |
| else: | |
| return 'lightcoral' | |
| except Exception as e: | |
| return 'lightgray' | |
| # Function to find model names in the family tree | |
| def get_model_names(model, genealogy, found_models=None, visited_models=None): | |
| if found_models is None: | |
| found_models = set() | |
| if visited_models is None: | |
| visited_models = set() | |
| if model in visited_models: | |
| 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: | |
| pass | |
| return found_models | |
| # Function to create the family tree | |
| def create_family_tree(start_model): | |
| genealogy = defaultdict(list) | |
| get_model_names(start_model, genealogy) | |
| G = nx.DiGraph() | |
| for parent, children in genealogy.items(): | |
| for child in children: | |
| G.add_edge(parent, child) | |
| max_depth = nx.dag_longest_path_length(G) + 1 | |
| max_width = max_width_of_tree(G) + 1 | |
| height = max(8, 1.6 * max_depth) | |
| width = max(8, 6 * max_width) | |
| plt.figure(figsize=(width, height)) | |
| pos = graphviz_layout(G, prog="dot") | |
| node_colors = [get_license_color(node) for node in G.nodes()] | |
| clear_output() | |
| labels = {node: node.replace("/", "\n") for node in G.nodes()} | |
| nx.draw(G, pos, labels=labels, with_labels=True, node_color=node_colors, font_size=12, node_size=8_000, edge_color='black') | |
| 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) | |
| plt.show() | |
| create_family_tree(MODEL_ID) | |