Adding modeling.py file

modeling.py

@@ -223,8 +223,19 @@ class MoLM(PreTrainedModel):
             # # 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)
+            log_probs = self(idx_cond, date=date).combined_log_probs[:, -1, :]
+            #idx_next = torch.multinomial(probs, num_samples=1)
+            # Sample from the log probabilities
+            if temperature == 0:
+                # If temperature is 0, take the argmax (greedy sampling)
+                idx_next = torch.argmax(log_probs, dim=-1, keepdim=True)
+            else:
+                # Apply temperature scaling
+                scaled_log_probs = log_probs / temperature
+                # Convert log probabilities to probabilities
+                probs = torch.exp(scaled_log_probs)
+                # Sample from the distribution
+                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)
             # check if we hit the end of the sequence