CAG: Cache-Augmented Generation

def get_token_count(text, model_name="meta-llama/Meta-Llama-3-8B"):
  """
  Calculates the number of tokens in the given text for a specified LLaMA model.

  Args:
    text: The input text string.
    model_name: The name of the LLaMA model (default: "meta-llama/Meta-Llama-3-8B").

  Returns:
    The number of tokens in the text.
  """
  tokenizer = AutoTokenizer.from_pretrained(model_name)
  tokens = tokenizer.encode(text)
  return len(tokens)

def get_text_from_file(file_path):
  """
  Reads the entire contents of a text file.

  Args:
    file_path: The path to the text file.

  Returns:
    The contents of the file as a single string.
  """
  try:
    with open(file_path, 'r') as file:
      text_data = file.read()
    return text_data
  except FileNotFoundError:
    print(f"Error: File not found at '{file_path}'")
    return None


data_path = os.path.abspath("shakespeare.txt")
text_data = get_text_from_file(data_path)
token_count = get_token_count(text_data)
print(f'Number of tokens in the file {data_path}: {"{:,}".format(token_count)}')

def generate(model, input_ids: torch.Tensor, past_key_values, max_new_tokens: int = 50) -> torch.Tensor:
    """
    Generates a sequence of tokens using the given model.
    Args:
        model: The language model to use for generation.
        input_ids (torch.Tensor): The input token IDs.
        past_key_values: The past key values for the model's attention mechanism.
        max_new_tokens (int, optional): The maximum number of new tokens to generate. Defaults to 50.
    Returns:
        torch.Tensor: The generated token IDs, excluding the input tokens.
    """
    device = model.model.embed_tokens.weight.device
    origin_len = input_ids.shape[-1]
    input_ids = input_ids.to(device)
    output_ids = input_ids.clone()
    next_token = input_ids

    with torch.no_grad():
        for _ in range(max_new_tokens):
            out = model(input_ids=next_token, past_key_values=past_key_values, use_cache=True)
            logits = out.logits[:, -1, :]
            token = torch.argmax(logits, dim=-1, keepdim=True)
            output_ids = torch.cat([output_ids, token], dim=-1)
            past_key_values = out.past_key_values
            next_token = token.to(device)

            if model.config.eos_token_id is not None and token.item() == model.config.eos_token_id:
                break

    return output_ids[:, origin_len:]


def get_kv_cache(model, tokenizer, prompt: str) -> DynamicCache:
    """
    Generates a key-value cache for a given model and prompt.
    Args:
        model: The language model to use for generating the cache.
        tokenizer: The tokenizer associated with the model.
        prompt (str): The input prompt for which the cache is generated.
    Returns:
        DynamicCache: The generated key-value cache.
    """
    device = model.model.embed_tokens.weight.device
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    cache = DynamicCache()

    with torch.no_grad():
        _ = model(input_ids=input_ids, past_key_values=cache, use_cache=True)
    return cache


def clean_up(cache: DynamicCache, origin_len: int):
    """
    Trims the key_cache and value_cache tensors in the given DynamicCache object.

    Args:
        cache (DynamicCache): The cache object containing key_cache and value_cache tensors.
        origin_len (int): The length to which the tensors should be trimmed.

    Returns:
        None
    """
    for i in range(len(cache.key_cache)):
        cache.key_cache[i] = cache.key_cache[i][:, :, :origin_len, :]
        cache.value_cache[i] = cache.value_cache[i][:, :, :origin_len, :]

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# load a model, here I chose phi because it is small with a big context
model_name = "microsoft/Phi-3-mini-128k-instruct"
hf_token = os.getenv("HF_TOKEN")
tokenizer = AutoTokenizer.from_pretrained(
    model_name, token=hf_token, trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
    trust_remote_code=True,
    token=hf_token,
)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print(f"Loaded {model_name}.")

#load a file for testing
with open("genesis.txt", "r", encoding="utf-8") as f:
    doc_text = f.read()

system_prompt = f"""
<|system|>
You are an assistant who provides concise factual answers.
You strive to just answer the user's question.
<|user|>
Context:
{doc_text}
Question:
""".strip()

genesis_cache = get_kv_cache(model, tokenizer, system_prompt)
origin_len = genesis_cache.key_cache[0].shape[-2]
print("KV cache built.")


# query the cache
question1 = "Why did God create eve?"
clean_up(genesis_cache, origin_len)
input_ids_q1 = tokenizer(question1 + "\n", return_tensors="pt").input_ids.to(device)
gen_ids_q1 = generate(model, input_ids_q1, genesis_cache)
answer1 = tokenizer.decode(gen_ids_q1[0], skip_special_tokens=True)
print("Q1:", question1)
print("A1:", answer1)