LucasLima commited on
Commit
ce4972d
·
verified ·
1 Parent(s): e963250

Update src/display/utils.py

Browse files
Files changed (1) hide show
  1. src/display/utils.py +8 -7
src/display/utils.py CHANGED
@@ -123,8 +123,7 @@ class ModelDetails:
123
 
124
  class ModelType(Enum):
125
  PT = ModelDetails(name="Pre trained", symbol="🟢")
126
- FT = ModelDetails(name="Fine-tuned", symbol="🔶")
127
- IFT = ModelDetails(name="Instruction-tuned", symbol="🔶")
128
  RL = ModelDetails(name="Reinforcement Learning", symbol="🟦")
129
  Unknown = ModelDetails(name="", symbol="?")
130
 
@@ -133,11 +132,13 @@ class ModelType(Enum):
133
 
134
  @staticmethod
135
  def from_str(type_str):
136
- if "fine-tuned" in type_str.lower() and "instruction" not in type_str.lower():
137
- return ModelType.FT
138
- if "instruction-tuned" in type_str.lower():
139
- return ModelType.IFT
140
- if "pretrained" in type_str.lower() or "pre trained" in type_str.lower() or "🟢" in type_str or type_str == "PT":
 
 
141
  return ModelType.PT
142
  if "rl-tuned" in type_str.lower() or "reinforcement learning" in type_str.lower() or "🟦" in type_str or type_str == "RL":
143
  return ModelType.RL
 
123
 
124
  class ModelType(Enum):
125
  PT = ModelDetails(name="Pre trained", symbol="🟢")
126
+ SFT = ModelDetails(name="Supervised Finetuning", symbol="🔶")
 
127
  RL = ModelDetails(name="Reinforcement Learning", symbol="🟦")
128
  Unknown = ModelDetails(name="", symbol="?")
129
 
 
132
 
133
  @staticmethod
134
  def from_str(type_str):
135
+ if "fine-tuned" in type_str.lower() or \
136
+ "instruction-tuned" in type_str.lower() or \
137
+ "supervised finetuning" in type_str.lower() or \
138
+ "🔶" in type_str or \
139
+ type_str == "SFT" or type_str == "FT" or type_str == "IFT":
140
+ return ModelType.SFT
141
+ if "pretrained" in type_str.lower() or "pre trained" in type_str.lower() or "pré-treinado" in type_str.lower() or "🟢" in type_str or type_str == "PT":
142
  return ModelType.PT
143
  if "rl-tuned" in type_str.lower() or "reinforcement learning" in type_str.lower() or "🟦" in type_str or type_str == "RL":
144
  return ModelType.RL