copy in the backend code

custom_llm.py

@@ -0,0 +1,170 @@
+import argparse
+import os
+import time
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.testclient import TestClient
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from custom_llm_inference import get_highlights_inner, get_next_token_predictions_inner
+
+ml_models = {}
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--gpu", action="store_true", help="Enable GPU usage")
+args = parser.parse_args()
+
+USE_GPU = args.gpu
+
+if not USE_GPU:
+    print("Running without GPU. To enable GPU, run with the --gpu flag.")
+
+@asynccontextmanager
+async def models_lifespan(app: FastAPI):
+
+    #model_name = 'google/gemma-1.1-7b-it'
+    #model_name = 'google/gemma-1.1-2b-it'
+    model_name = 'google/gemma-2-9b-it'
+
+    dtype = torch.bfloat16 if USE_GPU else torch.float16
+
+    ml_models["llm"] = llm = {
+        'tokenizer': AutoTokenizer.from_pretrained(model_name),
+        'model': AutoModelForCausalLM.from_pretrained(model_name, device_map="auto" if USE_GPU else "cpu", torch_dtype=dtype)
+    }
+    print("Loaded llm with device map:")
+    print(llm['model'].hf_device_map)
+
+    # Print timing info for each endpoint
+    print("\nRunning endpoint tests...")
+    
+    test_doc = "This is a test document that needs to be revised for clarity and conciseness."
+    test_prompt = "Make this more clear and concise."
+    
+    client = TestClient(app)
+    
+    start = time.time()
+    response = client.get("/api/highlights", 
+        params={"doc": test_doc, "prompt": test_prompt})
+    print(f"Highlights endpoint: {time.time() - start:.2f}s")
+    
+    start = time.time()
+    response = client.get("/api/next_token",
+        params={"original_doc": test_doc, "prompt": test_prompt, "doc_in_progress": "This is"})
+    print(f"Next token endpoint: {time.time() - start:.2f}s")
+    
+    start = time.time()
+    response = client.get("/api/gen_revisions",
+        params={"doc": test_doc, "prompt": test_prompt, "n": 1})
+    print(f"Gen revisions endpoint: {time.time() - start:.2f}s")
+
+    yield
+
+    # Release resources on exit
+    ml_models.clear()
+
+DEBUG = os.getenv("DEBUG") or False
+PORT = int(os.getenv("PORT") or "19570")
+
+app = FastAPI(lifespan=models_lifespan)
+
+origins = [
+    "*",
+]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.get("/api/highlights")
+def get_highlights(doc: str, prompt: Optional[str] = None, updated_doc: Optional[str] = '', k: Optional[int] = 5):
+    ''' Example of using this in JavaScript:
+    
+    let url = new URL('http://localhost:8000/api/highlights')
+    url.searchParams.append('doc', 'This is a test document. It is a test document because it is a test document.')
+    url.searchParams.append('prompt', 'Rewrite this document to be more concise.')
+    url.searchParams.append('updated_doc', 'This is a test document.')
+    let response = await fetch(url)
+    '''
+
+    llm = ml_models['llm']
+    model = llm['model']
+    tokenizer = llm['tokenizer']
+
+    if prompt is None:
+        prompt = "Rewrite this document to be more concise."
+
+    highlights = get_highlights_inner(model, tokenizer, doc, prompt, updated_doc, k)
+
+    return {'highlights': highlights}
+
+
+@app.get('/api/next_token')
+def get_next_token_predictions(original_doc: str,
+                               prompt: str,
+                               doc_in_progress: str,
+                               k: Optional[int] = 5):
+
+
+    model = ml_models['llm']['model']
+    tokenizer = ml_models['llm']['tokenizer']
+
+    decoded_next_tokens, next_token_logits = get_next_token_predictions_inner(
+        model, tokenizer, original_doc, prompt, doc_in_progress, k) 
+
+    return {
+        'next_tokens': decoded_next_tokens
+    }
+
+
+@app.get('/api/gen_revisions')
+def gen_revisions(
+        prompt: str,
+        doc: str,
+        n: Optional[int] = 5):
+
+
+    model = ml_models['llm']['model']
+    tokenizer = ml_models['llm']['tokenizer']
+
+    messages = [
+        {
+            "role": "user",
+            "content": f"{prompt}\n\n{doc}",
+        },
+    ]
+    tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
+
+    generations = model.generate(
+        tokenized_chat, num_return_sequences=n,
+        max_length=1024, do_sample=True, top_k=50, top_p=0.95, temperature=0.5,
+        return_dict_in_generate=True, output_scores=True)
+    generated_docs = tokenizer.batch_decode(generations.sequences, skip_special_tokens=True)
+    #print(generations.scores)
+
+    # Remove prompt text. see https://github.com/huggingface/transformers/blob/v4.46.2/src/transformers/pipelines/text_generation.py#L37
+    prompt_length = len(
+        tokenizer.decode(
+            tokenized_chat[0],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=True,
+    ))
+
+    return {
+        'revised_docs': [dict(doc_text=doc[prompt_length:]) for doc in generated_docs]
+    }
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="localhost", port=PORT)

custom_llm_inference.py

@@ -0,0 +1,193 @@
+import torch
+from transformers.cache_utils import DynamicCache
+
+
+def get_tokenized_chat(tokenizer, prompt, doc):
+    messages = [
+        {
+            "role": "user",
+            "content": f"{prompt}\n\n{doc}",
+        },
+    ]
+    tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")[0]
+    return tokenized_chat
+
+
+def tokenize_doc_in_progress(tokenizer, doc_in_progress):
+    if len(doc_in_progress) == 0:
+        # Some tokenizers give tensors of the wrong dtype if the input is empty
+        return torch.empty(0, dtype=torch.int64)
+
+    doc_in_progress_ids = tokenizer(
+        doc_in_progress, return_tensors='pt')['input_ids'][0]
+
+    # strip the first token, the "beginning of document" token
+    # TODO: make this robust to switching models
+    # since some models will use different special tokens
+    doc_in_progress_ids = doc_in_progress_ids[1:]
+    return doc_in_progress_ids
+
+
+def get_highlights_inner(model, tokenizer, doc, prompt, updated_doc, k):
+    tokenized_chat = get_tokenized_chat(tokenizer, prompt, doc)
+    assert len(tokenized_chat.shape) == 1
+
+    if updated_doc is None or len(updated_doc.strip()) == 0:
+        updated_doc = doc
+    updated_doc_ids = tokenize_doc_in_progress(tokenizer, updated_doc)
+
+    joined_ids = torch.cat([tokenized_chat, updated_doc_ids])
+    # Call the model
+    with torch.no_grad():
+        logits = model(joined_ids[None].to(model.device)).logits[0].cpu()
+    
+    highlights = []
+    length_so_far = 0
+    for idx in range(len(tokenized_chat), len(joined_ids)):
+        probs = logits[idx - 1].softmax(dim=-1)
+        token_id = joined_ids[idx]
+        token = tokenizer.decode(token_id)
+        token_loss = -probs[token_id].log().item()
+        topk_tokens = probs.topk(k).indices.cpu().numpy().tolist()
+        topk_tokens_decoded = tokenizer.batch_decode(topk_tokens, skip_special_tokens=True)
+        highlights.append(dict(
+            start=length_so_far,
+            end=length_so_far + len(token),
+            token=token,
+            token_loss=token_loss,
+            most_likely_token=topk_tokens_decoded[0],
+            topk_tokens=topk_tokens_decoded,
+        ))
+        length_so_far += len(token)
+    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
+
+    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
+
+    # 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))
+
+    # 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".
+    next_tokens_as_batch = branch_tokens.unsqueeze(1)
+    assert next_tokens_as_batch.shape == (k, 1)
+
+    position_id_for_final_token = joined_ids.shape[0]
+    cache_position = torch.full((1,), position_id_for_final_token, dtype=int, device=device)
+    with torch.no_grad():
+        model_outs = model(
+            next_tokens_as_batch,
+            past_key_values=past_key_values,
+            output_hidden_states=True,
+            use_cache=True,
+            # the cache surprisingly doesn't know the position of the last token
+            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 == (k, vocab_size), f"{next_token_logits.shape=}, {k=}, {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,)
+    lookahead_sequences = torch.cat([
+        branch_tokens.unsqueeze(1),
+        most_likely_token_ids.unsqueeze(1)
+    ], dim=1)
+    assert lookahead_sequences.shape == (k, 2)
+
+    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(
+        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)
+
+    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, 
+
+
+def get_next_token_predictions_slow(
+        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)
+
+    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
+
+    with torch.no_grad():
+        model_outs = model(hypotheses, output_hidden_states=True)
+
+    next_token_logits = model_outs.logits[0, -1]
+    branch_tokens = next_token_logits.topk(k).indices
+
+    # Slow mode: concat the branch tokens to the hypotheses.
+    # Then call the model on the full sequence.
+    # This is slow because the beginning of the sequence is re-processed each time.
+
+    hypotheses_with_next_tokens = torch.cat([
+        torch.repeat_interleave(hypotheses, k, dim=0),
+        branch_tokens.unsqueeze(1)
+    ], dim=1)
+    assert hypotheses_with_next_tokens.shape == (k, len(joined_ids) + 1)
+
+    with torch.no_grad():
+        model_outs = model(hypotheses_with_next_tokens)
+    
+    # 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 == (k, vocab_size), f"{next_token_logits.shape=}, {k=}, {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,)
+    lookahead_sequences = torch.cat([
+        branch_tokens.unsqueeze(1),
+        most_likely_token_ids.unsqueeze(1)
+    ], dim=1)
+    assert lookahead_sequences.shape == (k, 2)
+
+    decoded_next_tokens = tokenizer.batch_decode(lookahead_sequences, skip_special_tokens=True)
+    return decoded_next_tokens, next_token_logits