fix weight class

app.py

@@ -18,7 +18,6 @@ from src.utils import (
     load_dataset_repo,
     process_model_name,
     process_model_type,
-    process_weight_class,
 )
 
 
@@ -78,6 +77,11 @@ def get_benchmark_df(benchmark="1xA100-80GB"):
     scores_df["merge_id"] = scores_df.weight_class + "_" + scores_df.model_type
     merged_df = bench_df.merge(scores_df, on="merge_id")
 
+    # fix some weight classes
+    merged_df["weight_class"] = merged_df["weight_class"].apply(
+        lambda x: TRUE_WEIGHT_CLASSES[x] if x in TRUE_WEIGHT_CLASSES else x
+    )
+
     # add optimizations
     merged_df["optimizations"] = merged_df[
         ["backend.bettertransformer", "backend.load_in_8bit", "backend.load_in_4bit"]

src/utils.py

@@ -41,6 +41,7 @@ LLM_MODEL_TYPES = {
     "RefinedWebModel": "Falcon 🦅",
     "RefinedWeb": "Falcon 🦅",
     "baichuan": "Baichuan 🌊",
+    "bloom": "Bloom 🌸",
     "llama": "LLaMA 🦙",
     "gpt_neox": "GPT-NeoX",
     "gpt_neo": "GPT-Neo",
@@ -68,15 +69,3 @@ def process_model_type(model_type):
         return LLM_MODEL_TYPES[model_type]
     else:
         return model_type
-
-
-def process_weight_class(num):
-    if num < 1000:
-        return str(int(num))
-    elif num < 1000000:
-        return str(int(num / 1000)) + "K"
-    elif num < 1000000000:
-        return str(int(num / 1000000)) + "M"
-    elif num < 1000000000000:
-        return str(int(num / 1000000000)) + "B"
-    return None