Skip to content

Commit

Permalink
Update gen_readme.py
Browse files Browse the repository at this point in the history
  • Loading branch information
rickzx committed Jun 19, 2024
1 parent 889dc7d commit 1bc4d50
Showing 1 changed file with 42 additions and 33 deletions.
75 changes: 42 additions & 33 deletions gen_readme.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,66 +4,75 @@
from collections import defaultdict
from jinja2 import Environment, FileSystemLoader

root_dir = 'bentoml/bentos'
root_dir = "bentoml/bentos"
grouped_data = defaultdict(list)

model_display_names = {
'llama3': 'Llama-3',
'phi3': 'Phi-3',
'mistral': 'Mistral',
'qwen2': 'Qwen-2',
'gemma': 'Gemma',
'llama2': 'Llama-2',
"llama3": "Llama-3",
"phi3": "Phi-3",
"mistral": "Mistral",
"qwen2": "Qwen-2",
"gemma": "Gemma",
"llama2": "Llama-2",
}

model_priority = {
'llama3': 1,
'phi3': 2,
'mistral': 3,
'qwen2': 4,
'gemma': 5,
'llama2': 6,
"llama3": 1,
"phi3": 2,
"mistral": 3,
"qwen2": 4,
"gemma": 5,
"llama2": 6,
}

yaml_files = glob.glob(os.path.join(root_dir, '**/bento.yaml'), recursive=True)
yaml_files = glob.glob(os.path.join(root_dir, "**/bento.yaml"), recursive=True)

# Read each bento.yaml file and group data by "name" field
for yaml_file in yaml_files:
with open(yaml_file, 'r') as f:
with open(yaml_file, "r") as f:
data = yaml.safe_load(f)
# Extract the HF model ID from routes.input.model.default
for route in data.get('schema', {}).get('routes', []):
for prop, details in route.get('input', {}).get('properties', {}).items():
if prop == 'model' and 'default' in details:
data['hf_model'] = details['default']
for route in data.get("schema", {}).get("routes", []):
for prop, details in route.get("input", {}).get("properties", {}).items():
if prop == "model" and "default" in details:
data["hf_model"] = details["default"]
break
# Append data to grouped_data
if 'name' in data and 'version' in data and 'hf_model' in data:
grouped_data[data['name']].append({
'name': data['name'],
'version': data['version'],
'hf_model': data['hf_model']
})
if "name" in data and "version" in data and "hf_model" in data:
grouped_data[data["name"]].append(
{
"name": data["name"],
"version": data["version"],
"hf_model": data["hf_model"],
}
)

# Deduplicate entries by converting lists to sets of tuples and back to lists
# Deduplicate entries
for name in grouped_data:
seen = set()
deduped_items = []
for item in grouped_data[name]:
item_tuple = (item['name'], item['version'], item['hf_model'])
item_tuple = (item["name"], item["version"], item["hf_model"])
if item_tuple not in seen:
seen.add(item_tuple)
deduped_items.append(item)
grouped_data[name] = deduped_items

# Sort the items within each group by version
for name in grouped_data:
grouped_data[name] = sorted(grouped_data[name], key=lambda x: x["version"])

# Sort the model names by priority and transform grouped_data into a list
sorted_grouped_data = sorted(grouped_data.items(), key=lambda x: model_priority.get(x[0], float('inf')))
sorted_grouped_data = sorted(
grouped_data.items(), key=lambda x: model_priority.get(x[0], float("inf"))
)

# Set up Jinja environment and load the template
env = Environment(loader=FileSystemLoader('.'))
template = env.get_template('readme_md.tpl')
env = Environment(loader=FileSystemLoader("."))
template = env.get_template("readme_md.tpl")

readme_content = template.render(grouped_data=sorted_grouped_data, model_display_names=model_display_names)
readme_content = template.render(
grouped_data=sorted_grouped_data, model_display_names=model_display_names
)

with open('README.md', 'w') as readme_file:
with open("README.md", "w") as readme_file:
readme_file.write(readme_content)

0 comments on commit 1bc4d50

Please sign in to comment.