Upload PAL_B_RM_opt

README.md

@@ -1,12 +1,12 @@
 ---
-library_name: transformers
-license: mit
+base_model:
+- facebook/opt-350m
 datasets:
 - CarperAI/openai_summarize_tldr
 language:
 - en
-base_model:
-- facebook/opt-350m
+library_name: transformers
+license: mit
 ---
 
 # Model Card for Model ID

itemLearner.py

@@ -27,6 +27,8 @@ class ItemLearner(nn.Module):
         '''
         input_ids = x['input_ids']
         attention_mask = x['attention_mask']
+        # logger.critical(f"ItemLearner: {input_ids=}")
+        # logger.critical(f"ItemLearner: {attention_mask=}")
 
         if rm_cached is None:
             llm_res = self.llm(
@@ -37,11 +39,12 @@ class ItemLearner(nn.Module):
             llm_res = self.llm(
                 input_ids=input_ids[:, -1:], # attention_mask=attention_mask,
                 past_key_values=rm_cached["item_learner"],
-                use_cache=False
+                use_cache=True
             )
             rm_cached["item_learner"] = llm_res.past_key_values
 
         embeds = llm_res.last_hidden_state
+        # logger.critical(f"ItemLearner: {embeds=}")
         # embeds shape: (bs, seq_len, hidden_size)
         shape = embeds.shape
         embeds = embeds.view(-1, shape[-1]) # (bs*seq_len, hidden_size)

learner.py

@@ -113,16 +113,18 @@ class PrefLearner(BasePrefLearner):   # <f(x),f(u)>
 
     def forward(self, x, rm_cached=None):
         assert self.uid is not None, "Please specify the user id first by calling specify_user_ids() to personalize the reward model"
-        items, prompt = x
+        prompt, items = x
         if rm_cached is None:
             items_prime, prompt_prime = self.map_to_pref_embedding_space((self.uid, prompt, items))
         else:
             items_prime, prompt_prime, rm_cached = self.map_to_pref_embedding_space((self.uid, prompt, items), rm_cached)
-        logger.info(f"{items_prime[0]=}")
-        logger.info(f"{prompt_prime[0]=}")
-        logger.info(f"{items_prime.shape=}")
-        logger.info(f"{prompt_prime.shape=}")
+        # logger.critical(f"{items_prime[0]=}")
+        # 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!
             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)
@@ -131,8 +133,8 @@ class PrefLearner(BasePrefLearner):   # <f(x),f(u)>
             items_last_prime = items_last_prime / torch.norm(items_last_prime, dim=-1, keepdim=True)
             logit_scale = self.logit_scale.exp()
             clamped_logit_scale = torch.clamp(logit_scale, max=100)
-            logger.info(f"{prompt_last_prime.shape=}")
-            logger.info(f"{items_last_prime.shape=}")
+            # logger.critical(f"{prompt_last_prime.shape=}")
+            # logger.critical(f"{items_last_prime.shape=}")
             sim_score = (prompt_last_prime * items_last_prime).sum(dim=-1) * clamped_logit_scale   # (bs, max_token_length)
             if rm_cached is None:
                 return sim_score

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:770877c170b8b51c6e6555de658213f0a6a1fca5c74370f1b8fed47cf6411bac
+oid sha256:8c5e8e5083c6c333b9ba3284e989dab95ef00f70e1b97770df255acebada4388
 size 1334487698

userLearner.py

@@ -92,6 +92,11 @@ class UserLearner(nn.Module):
 
         # embeds shape: (bs, seq_len, hid_dim)
         shape = embeds.shape
+        # only last hidden state start
+        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
+        # logger.critical("using only last hidden state of prompt tokens")
         embeds = embeds.view(-1, shape[-1])  # (bs*seq_len, hid_dim)
         # g(embeds) shape: (bs*seq_len, hid_dim) -> (bs*seq_len, pref_dim)
         logits = torch.stack([g(embeds).view(shape[0], shape[1], -1) for g in self.projectors.values()],dim=1)
@@ -118,8 +123,8 @@ class UserLearner(nn.Module):
         # assert sum(mix_weight) == 1
         # w = self.softmax(mix_weight.repeat(bs, 1))
         # w = mix_weight.repeat(bs, 1)
-        logger.info(f"{w=}")
-        logger.info(f"{w.shape=}")
+        # logger.info(f"{w=}")
+        # logger.info(f"{w.shape=}")
         w = w.unsqueeze(-1).unsqueeze(-1)
         y_hat = (w * prompt_logits).sum(dim=1)
         self.tmp_store_user_ideal_points = y_hat