bugfixes

w2s/loss.py

@@ -18,7 +18,13 @@ def confidence_window_loss(
 
  target = torch.stack([1.0 - labels, labels], dim=1)
 
- return torch.nn.functional.cross_entropy(logits[uncertain], target[uncertain])
+ loss = torch.nn.functional.cross_entropy(
+ logits[uncertain], 
+ target[uncertain], 
+ reduction="sum"
+ )
+
+ return loss / logits.shape[0]
 
 
 def log_confidence_loss(

w2s/sft.py

@@ -54,7 +54,7 @@ def compute_loss(self, model, inputs, return_outputs=False):
  loss = confidence_window_loss(
  outputs.logits,
  labels,
- radius=self.loss_cfg.radius,
+ radius=(self.loss_cfg.radius if self.transfer else 0.51),
  )
  else:
  raise ValueError(f"Unknown loss function: {self.loss_name}")
@@ -113,8 +113,8 @@ def compute_metrics(eval_pred):
  )
 
  trainer = CustomLossTrainer(
- loss_name=cfg.loss_name,
- loss_cfg=cfg.loss,
+ loss_name=cfg["loss_name"],
+ loss_cfg=cfg["loss"],
  transfer=transfer,
  args=train_args,
  compute_metrics=compute_metrics,
@@ -140,7 +140,8 @@ def compute_metrics(eval_pred):
  with open(save_dir / "config.json", "w") as f:
  cfg["model"] = model_cfg.to_dict()
  cfg["train_args"] = train_args.to_dict()
- cfg["logconf_weight"] = logconf_weight
+ cfg["transfer"] = transfer
+ cfg["loss"] = cfg["loss"].to_dict()
  json.dump(cfg, f, indent=2)
  wandb.config.update(cfg)
 

w2s/sft_config.py

@@ -6,7 +6,9 @@
 
 @dataclass
 class LossConfig(Serializable):
- pass
+ def to_dict(self):
+ irrelevant_fields = []
+ return {k: v for k, v in vars(self).items() if k not in irrelevant_fields}
 
 @dataclass
 class LogConfidenceLossConfig(LossConfig):

w2s/utils.py

@@ -1,4 +1,5 @@
 from typing import Any, Type, TypeVar, cast
+from w2s.sft_config import LossConfig
 
 T = TypeVar("T")
 
@@ -10,6 +11,12 @@ def assert_type(typ: Type[T], obj: Any) -> T:
 
  return cast(typ, obj)
 
+NICKNAMES = {
+ "Qwen/Qwen1.5-0.5B": "Qw0.5",
+ "meta-llama/Meta-Llama-3-8B": "Ll8",
+ "./results": "rs",
+ "cosine": "c",
+}
 
 def get_config_foldername(config: dict) -> str:
  def shorten_key(key: str) -> str:
@@ -19,6 +26,8 @@ def shorten_value(value) -> str:
  if isinstance(value, bool):
  return "1" if value else "0"
  elif isinstance(value, str):
+ if value in NICKNAMES:
+ return NICKNAMES[value]
  value = value.split("/")[-1]
  if "_" in value:
  return "_".join(word[:4] for word in value.split("_"))
@@ -39,6 +48,8 @@ def flatten_dict(d: dict, parent_key: str = "", sep: str = "_") -> dict:
  new_key = parent_key + sep + k if parent_key else k
  if isinstance(v, dict):
  items.extend(flatten_dict(v, new_key, sep=sep).items())
+ elif isinstance(v, LossConfig):
+ items.extend(flatten_dict(v.to_dict(), new_key, sep=sep).items())
  else:
  items.append((new_key, v))
  return dict(items)