Adding modeling.py file

modeling.py

@@ -212,17 +212,18 @@ class MoLM(PreTrainedModel):
                 if idx.size(1) <= self.config.sequence_length
                 else idx[:, -self.config.sequence_length :]
             )
-            # forward the model to get the logits for the index in the sequence
-            logits = self(idx_cond, date, get_logits=True).logits
-            # pluck the logits at the final step and scale by desired temperature
-            logits = logits[:, -1, :] / temperature
-            # optionally crop the logits to only the top k options
-            if top_k is not None:
-                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
-                logits[logits < v[:, [-1]]] = -float("Inf")
-            # apply softmax to convert logits to (normalized) probabilities
-            probs = F.softmax(logits, dim=-1)
-            # sample from the distribution
+            # # forward the model to get the logits for the index in the sequence
+            # logits = self(idx_cond, date, get_logits=True).logits
+            # # pluck the logits at the final step and scale by desired temperature
+            # logits = logits[:, -1, :] / temperature
+            # # optionally crop the logits to only the top k options
+            # if top_k is not None:
+            #     v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+            #     logits[logits < v[:, [-1]]] = -float("Inf")
+            # # apply softmax to convert logits to (normalized) probabilities
+            # probs = F.softmax(logits, dim=-1)
+            # # sample from the distribution
+            probs = self(idx_cond, date=date).combined_log_probs[:, -1, :]
             idx_next = torch.multinomial(probs, num_samples=1)
             # append sampled index to the running sequence and continue
             idx = torch.cat((idx, idx_next), dim=1)