Initial commit

custom_generate/generate.py

@@ -103,9 +103,9 @@ def ancestral_sampling(model_kwargs, model, eos_token_ids, pad_token_id, bos_tok
     active_seqs = input_ids.new_ones((batch_size, 1), dtype=torch.bool)
     lens = torch.full((batch_size,), max_prompts_len, dtype=torch.long, device=input_ids.device)
     # Modified log probabilities of the sequences
-    scores = torch.zeros((batch_size, max_new_tokens), dtype=torch.float32)
+    scores = torch.zeros((batch_size, max_new_tokens), dtype=model.dtype)
     # Unfiltered sequence log probabilities (T=1, no sampling modifications)
-    logps = torch.zeros((batch_size, max_new_tokens), dtype=torch.float32)
+    logps = torch.zeros((batch_size, max_new_tokens), dtype=model.dtype)
 
     for i in range(max_new_tokens):
         # Get the next token probabilities and update the KV cache
@@ -153,6 +153,7 @@ def generate(model, **kwargs):
     """
     generation_config = model.generation_config
     max_new_tokens = kwargs.get('max_new_tokens', generation_config.max_new_tokens)
+    max_new_tokens = 512 if max_new_tokens is None else max_new_tokens
     do_sample = kwargs.get('do_sample', True)
     eos_token_ids = kwargs.get('eos_token_ids', generation_config.eos_token_id)
     if eos_token_ids is None: