Upload PAL_B_RM_opt

learner.py

@@ -122,9 +122,8 @@ class PrefLearner(BasePrefLearner):   # <f(x),f(u)>
         # logger.critical(f"{prompt_prime[0]=}")
         # logger.critical(f"{items_prime.shape=}")
         # logger.critical(f"{prompt_prime.shape=}")
-        # FIXME: bug exist here 
         if self.pref_learner_type == 'angle':
-            # FIXME: do the cumulative evaluation!
+            # NOTICE: here we implement the "last token only" version of PAL-B
             prompt_last_prime = prompt_prime[:, -1, :]
             prompt_last_prime = prompt_last_prime.unsqueeze(1)
             prompt_last_prime = prompt_last_prime / torch.norm(prompt_last_prime, dim=-1, keepdim=True)

userLearner.py

@@ -92,7 +92,7 @@ class UserLearner(nn.Module):
 
         # embeds shape: (bs, seq_len, hid_dim)
         shape = embeds.shape
-        # only last hidden state start
+        # only last hidden state start (only use the last token of the prompt)
         embeds = embeds[:, -1, :]  # (bs, seq_len, hid_dim) -> (bs, hid_dim)
         embeds = embeds.unsqueeze(1).repeat(1, shape[1], 1)  # (bs, hid_dim) -> (bs, seq_len, hid_dim)
         # only last hidden state end