Looks like I never committed these improvements to the backend.

custom_llm_inference.py

@@ -63,37 +63,29 @@ def get_highlights_inner(model, tokenizer, doc, prompt, updated_doc, k):
     return highlights
 
 
-
-def get_next_token_predictions_inner(
-        model, tokenizer, original_doc, prompt, doc_in_progress, k):
-
-    tokenized_chat = get_tokenized_chat(tokenizer, prompt, original_doc)
-    doc_in_progress_ids = tokenize_doc_in_progress(tokenizer, doc_in_progress)
-
-    device = model.device
-
-    joined_ids = torch.cat([tokenized_chat, doc_in_progress_ids])
-    hypotheses = joined_ids[None].to(model.device)
-
-    # For each of the k next tokens, generate most-likely next tokens and append back on until we
-    # reach a token with a space
-
+def get_lookahead_sequences(model, tokenizer, hypotheses, n_branch_tokens, device):
+    """
+    For each of the n_branch_tokens next tokens, generate most-likely next tokens and append back on.
+    """
+    assert len(hypotheses.shape) == 2
+    assert hypotheses.shape[0] == 1
+    n_tokens_so_far = hypotheses.shape[1]
     past_key_values = DynamicCache()
 
     with torch.no_grad():
         model_outs_onestep = model(hypotheses, output_hidden_states=True, past_key_values=past_key_values)
 
-    branch_tokens = model_outs_onestep.logits[0, -1].topk(k).indices
+    branch_tokens = model_outs_onestep.logits[0, -1].topk(n_branch_tokens).indices
 
-    # split the cache into k reps. We pretend we're doing a "Beam search"...
-    past_key_values.reorder_cache(torch.zeros((k,), dtype=torch.long, device=device))
+    # split the cache into n_branch_tokens reps. We pretend we're doing a "Beam search"...
+    past_key_values.reorder_cache(torch.zeros((n_branch_tokens,), dtype=torch.long, device=device))
 
     # Now call the model again, passing the kv cache, so we can continue generating.
-    # Each of the k next tokens will be considered as one sequence in a "batch".
+    # Each of the n_branch_tokens next tokens will be considered as one sequence in a "batch".
     next_tokens_as_batch = branch_tokens.unsqueeze(1)
-    assert next_tokens_as_batch.shape == (k, 1)
+    assert next_tokens_as_batch.shape == (n_branch_tokens, 1)
 
-    position_id_for_final_token = joined_ids.shape[0]
+    position_id_for_final_token = n_tokens_so_far
     cache_position = torch.full((1,), position_id_for_final_token, dtype=int, device=device)
     with torch.no_grad():
         model_outs = model(
@@ -105,44 +97,52 @@ def get_next_token_predictions_inner(
             cache_position=cache_position
         )
     
-    # Grab the single most likely token from each of the k sequences
+    # Grab the single most likely token from each of the n_branch_tokens sequences
     next_token_logits = model_outs.logits[:, -1]
     vocab_size = model.config.vocab_size
-    assert next_token_logits.shape == (k, vocab_size), f"{next_token_logits.shape=}, {k=}, {vocab_size=}"
+    assert next_token_logits.shape == (n_branch_tokens, vocab_size), f"{next_token_logits.shape=}, {n_branch_tokens=}, {vocab_size=}"
     most_likely_token_ids = next_token_logits.argmax(dim=-1)
 
     # Stick them at the end of the branch tokens.
-    assert most_likely_token_ids.shape == (k,)
+    assert most_likely_token_ids.shape == (n_branch_tokens,)
     lookahead_sequences = torch.cat([
         branch_tokens.unsqueeze(1),
         most_likely_token_ids.unsqueeze(1)
     ], dim=1)
-    assert lookahead_sequences.shape == (k, 2)
+    assert lookahead_sequences.shape == (n_branch_tokens, 2)
+    return lookahead_sequences, next_token_logits
 
-    decoded_next_tokens = tokenizer.batch_decode(lookahead_sequences, skip_special_tokens=True)
-    return decoded_next_tokens, next_token_logits
 
-def get_next_token_predictions_generate(
+def get_next_token_predictions_inner(
         model, tokenizer, original_doc, prompt, doc_in_progress, k):
 
     tokenized_chat = get_tokenized_chat(tokenizer, prompt, original_doc)
     doc_in_progress_ids = tokenize_doc_in_progress(tokenizer, doc_in_progress)
 
+    device = model.device
+
     joined_ids = torch.cat([tokenized_chat, doc_in_progress_ids])
-    context_without_special_tokens = tokenizer.batch_decode(joined_ids, skip_special_tokens=True)
-    prefix_length = len(context_without_special_tokens)
     hypotheses = joined_ids[None].to(model.device)
 
-    generation_output = model.generate(
-        hypotheses,
-        return_dict_in_generate=True,
-        output_scores=True,
-        num_beams=5, num_beam_groups=5, max_new_tokens=10, do_sample=False, diversity_penalty=1e5, top_k=None, num_return_sequences=5)#, token_healing=True, tokenizer=tokenizer)
-    sequences = [
-        decoded[prefix_length:]
-        for decoded in tokenizer.batch_decode(generation_output.sequences, skip_special_tokens=True)
-    ]
-    return sequences, 
+    # Alternative approach: chat templates
+    tokenized_chat = tokenizer.apply_chat_template([
+        {"role": "user", "content": f"{prompt}\n\n{original_doc}"},
+        {"role": "assistant", "content": doc_in_progress}
+    ], tokenize=True, return_tensors="pt", continue_final_message=True).to(model.device)
+
+    # Compare them
+    if tokenized_chat.shape == hypotheses.shape and torch.all(tokenized_chat == hypotheses):
+        print("Tokenized chat and hypotheses match")
+    else:
+        print("FAIL: Tokenized chat and hypotheses do not match!")
+        print(f"{tokenized_chat=}")
+        print(f"{hypotheses=}")
+
+    lookahead_sequences, next_token_logits = get_lookahead_sequences(
+        model, tokenizer, hypotheses, k, device)
+
+    decoded_next_tokens = tokenizer.batch_decode(lookahead_sequences, skip_special_tokens=True)
+    return decoded_next_tokens, next_token_logits
 
 
 def get_next_token_predictions_slow(
@@ -196,67 +196,14 @@ def get_next_token_predictions_slow(
 
 
 def continue_messages_inner(model, tokenizer, messages, n_branch_tokens, n_future_tokens):
+    # Note: we're ignoring n_future_tokens right now since the old implementation was buggy.
     device = model.device
 
-    final_message_is_assistant = messages[-1]['role'] == "assistant"
-    print(f"final_message_is_assistant: {final_message_is_assistant}")
-    # if final_message_is_assistant:
-    #     tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, continue_final_message=True, return_tensors="pt").to(model.device)
-    # else:
-    #     tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
     tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, return_tensors="pt", continue_final_message=True).to(model.device)
-
     print(tokenizer.batch_decode(tokenized_chat, skip_special_tokens=False))
 
-    # This fails with
-    # RuntimeError: Index put requires the source and destination dtypes match, got BFloat16 for the destination and Float for the source.
-    # generations = model.generate(
-    #     tokenized_chat,
-    #     num_return_sequences=n_branch_tokens,
-    #     num_beam_groups=n_branch_tokens, num_beams=n_branch_tokens,
-    #     do_sample=False, max_new_tokens=n_future_tokens, diversity_penalty=1e5, top_k=None,
-    #     return_dict_in_generate=True, output_scores=True)
-
-    # Instead, we'll do this in two steps:
-    # 1. Get the next token predictions for the k most likely continuations
-    from transformers.cache_utils import DynamicCache
-    past_key_values = DynamicCache()
-    with torch.no_grad():
-        model_outs = model(
-            tokenized_chat,
-            past_key_values=past_key_values,
-            output_hidden_states=True,
-            use_cache=True,
-        )
-        branch_tokens = model_outs.logits[0, -1].topk(n_branch_tokens).indices
-    
-    hypotheses = branch_tokens.unsqueeze(1)
-    # Branch off the k most likely continuations
-    past_key_values.reorder_cache(torch.zeros((n_branch_tokens,), dtype=torch.long, device=device))
+    lookahead_sequences, next_token_logits = get_lookahead_sequences(
+        model, tokenizer, tokenized_chat, n_branch_tokens, device)
 
-    # 2. Generate the next n_future_tokens for each branch
-    for i in range(n_future_tokens):
-        position_id_for_final_token = tokenized_chat.shape[0] + i
-        cache_position = torch.full((1,), position_id_for_final_token, dtype=int, device=device)
-        final_token_ids = hypotheses[:, -1:]
-        with torch.no_grad():
-            model_outs = model(
-                final_token_ids,
-                past_key_values=past_key_values,
-                output_hidden_states=True,
-                use_cache=True,
-                cache_position=cache_position
-            )
-
-        # Grab the single most likely token from each of the k sequences
-        next_token_logits = model_outs.logits[:, -1]
-        vocab_size = model.config.vocab_size
-        assert next_token_logits.shape == (n_branch_tokens, vocab_size), f"{next_token_logits.shape=}, {n_branch_tokens=}, {vocab_size=}"
-        most_likely_token_ids = next_token_logits.argmax(dim=-1)
-        hypotheses = torch.cat([
-            hypotheses,
-            most_likely_token_ids.unsqueeze(1)
-        ], dim=1)
-
-    generated_docs = tokenizer.batch_decode(hypotheses, skip_special_tokens=True)
+    generated_docs = tokenizer.batch_decode(lookahead_sequences, skip_special_tokens=True)
     return generated_docs