Update modeling_gator.py

modeling_gator.py

@@ -100,13 +100,10 @@ class GatorForCausalLM(PreTrainedModel):
         self.gator = GatorModel(config)
         self.post_init()
 
-    def forward(self, input_ids):
-        return self.gator(input_ids)["logits"]
-
     @torch.no_grad()
-    def generate_one(self, input_ids, temperature=0.8, top_k=5):
-        logits = self.forward(input_ids)[:, -1, :] / temperature
+    def forward(self, input_ids, temperature=0.8, top_k=5):
+        logits = self.gator(input_ids)["logits"][:, -1, :] / temperature
         topk = torch.topk(logits, k=top_k, dim=-1)
         probs = torch.softmax(topk.values, dim=-1)
-        next_token = topk.indices[torch.multinomial(probs, num_samples=1)]
-        return next_token.item()
+        next_token = topk.indices.gather(-1, torch.multinomial(probs, 1))
+        return next_token.squeeze().item()