Update src/display/utils.py

src/display/utils.py

@@ -61,11 +61,33 @@ class ModelDetails:
     symbol: str = "" # emoji
 
 
+# class ModelType(Enum):
+#     PT = ModelDetails(name="pretrained", symbol="🟢")
+#     FT = ModelDetails(name="fine-tuned", symbol="🔶")
+#     IFT = ModelDetails(name="instruction-tuned", symbol="⭕")
+#     RL = ModelDetails(name="RL-tuned", symbol="🟦")
+#     Unknown = ModelDetails(name="", symbol="?")
+
+#     def to_str(self, separator=" "):
+#         return f"{self.value.symbol}{separator}{self.value.name}"
+
+#     @staticmethod
+#     def from_str(type):
+#         if "fine-tuned" in type or "🔶" in type:
+#             return ModelType.FT
+#         if "pretrained" in type or "🟢" in type:
+#             return ModelType.PT
+#         if "RL-tuned" in type or "🟦" in type:
+#             return ModelType.RL
+#         if "instruction-tuned" in type or "⭕" in type:
+#             return ModelType.IFT
+#         return ModelType.Unknown
 class ModelType(Enum):
     PT = ModelDetails(name="pretrained", symbol="🟢")
-    FT = ModelDetails(name="fine-tuned", symbol="🔶")
-    IFT = ModelDetails(name="instruction-tuned", symbol="⭕")
-    RL = ModelDetails(name="RL-tuned", symbol="🟦")
+    CPT = ModelDetails(name="continuously pretrained", symbol="🟩")
+    FT = ModelDetails(name="fine-tuned on domain-specific datasets", symbol="🔶")
+    chat = ModelDetails(name="chat models (RLHF, DPO, IFT, ...)", symbol="💬")
+    merges = ModelDetails(name="base merges and moerges", symbol="🤝")
     Unknown = ModelDetails(name="", symbol="?")
 
     def to_str(self, separator=" "):
@@ -75,41 +97,44 @@ class ModelType(Enum):
     def from_str(type):
         if "fine-tuned" in type or "🔶" in type:
             return ModelType.FT
+        if "continously pretrained" in type or "🟩" in type:
+            return ModelType.CPT
         if "pretrained" in type or "🟢" in type:
             return ModelType.PT
-        if "RL-tuned" in type or "🟦" in type:
-            return ModelType.RL
-        if "instruction-tuned" in type or "⭕" in type:
-            return ModelType.IFT
+        if any([k in type for k in ["instruction-tuned", "RL-tuned", "chat", "🟦", "⭕", "💬"]]):
+            return ModelType.chat
+        if "merge" in type or "🤝" in type:
+            return ModelType.merges
         return ModelType.Unknown
 
+
 class WeightType(Enum):
     Adapter = ModelDetails("Adapter")
     Original = ModelDetails("Original")
     Delta = ModelDetails("Delta")
 
 class Precision(Enum):
+    float32 = ModelDetails("float32")
     float16 = ModelDetails("float16")
     bfloat16 = ModelDetails("bfloat16")
-    float32 = ModelDetails("float32")
-    #qt_8bit = ModelDetails("8bit")
-    #qt_4bit = ModelDetails("4bit")
-    #qt_GPTQ = ModelDetails("GPTQ")
+    qt_8bit = ModelDetails("8bit")
+    qt_4bit = ModelDetails("4bit")
+    qt_GPTQ = ModelDetails("GPTQ")
     Unknown = ModelDetails("?")
 
     def from_str(precision):
+        if precision in ["float32"]:
+            return Precision.float32
         if precision in ["torch.float16", "float16"]:
             return Precision.float16
         if precision in ["torch.bfloat16", "bfloat16"]:
             return Precision.bfloat16
-        if precision in ["float32"]:
-            return Precision.float32
-        #if precision in ["8bit"]:
-        #    return Precision.qt_8bit
-        #if precision in ["4bit"]:
-        #    return Precision.qt_4bit
-        #if precision in ["GPTQ", "None"]:
-        #    return Precision.qt_GPTQ
+        if precision in ["8bit"]:
+           return Precision.qt_8bit
+        if precision in ["4bit"]:
+           return Precision.qt_4bit
+        if precision in ["GPTQ", "None"]:
+           return Precision.qt_GPTQ
         return Precision.Unknown
 
 # Column selection