Update src/display/utils.py

src/display/utils.py

@@ -62,25 +62,29 @@ class ModelDetails:
 
 
 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="?")
+    PT = ModelDetails(name="🟢 pretrained", 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="❓ other", 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:
+    def from_str(m_type):
+        if any([k for k in m_type if k in ["fine-tuned","🔶", "finetuned"]]):
             return ModelType.FT
-        if "pretrained" in type or "🟢" in type:
+        if "continuously pretrained" in m_type or "🟩" in m_type:
+            return ModelType.CPT
+        if "pretrained" in m_type or "🟢" in m_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 m_type for k in ["instruction-tuned", "RL-tuned", "chat", "🟦", "⭕", "💬"]]):
+            return ModelType.chat
+        if "merge" in m_type or "🤝" in m_type:
+            return ModelType.merges
         return ModelType.Unknown
 
 class WeightType(Enum):