Few-Shot + RAG on Llama

gistfile1.txt

# Download the Yelp Review Polarity CSV and unzip alongside this code under `yelp_review_polarity_csv`.
# Read more at https://joecooper.me/blog/ragshot/

# Config
n_train = 5600 # How much labeled data to use?
n_test = 38000 # How many test samples to run?
n_examples = 2 # How many examples per inference?
batch_size = 8 # Set according to your RAM

# Reszta
from os.path import exists
from sys import stderr
def load_csv(filename: str) -> list[tuple[str, str]]:
    if not exists(filename):
        print(f"{filename} not found!", file=stderr)
    with open(filename, 'r') as f:
        import csv
        r = csv.reader(f)
        return list((row[1], row[0]) for row in r)
train = load_csv('yelp_review_polarity_csv/train.csv')[:n_train]
def enumerate_batch(l: list, n: int):
    while len(l) > 0:
        b = l[:n]
        l = l[n:]
        yield b
import torch
from torch import tensor
from torch.utils.data import TensorDataset, DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch.nn.functional as F
device = 'cuda'
def load(s: str) -> tuple[tensor, tensor, tensor, tensor]:
    with open(s, 'rb') as f:
        return torch.load(f, map_location=device)
model_id = "meta-llama/Llama-3.2-1B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    use_fast=True)
tokenizer.pad_token = tokenizer.eos_token
dtype = torch.bfloat16
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=dtype,
    device_map=device)
pad_id = 2
def encode_fragment(s: str) -> torch.Tensor:
    b = tokenizer(s, return_tensors="pt", add_special_tokens=False)
    ids = b['input_ids'].squeeze(dim=0)
    return ids
model.eval()
def embed_batch(
    input_ids = torch.empty(1, 0),
    attention_mask = torch.empty(1, 0)) -> torch.Tensor:
    with torch.no_grad():
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True)
        last_hidden_state = outputs.hidden_states[-1]
        attention_mask_expanded = attention_mask.unsqueeze(-1)
        sum_embeddings = torch.sum(last_hidden_state * attention_mask_expanded, 1)
        sum_mask = torch.clamp(attention_mask_expanded.sum(1), min=1e-9) # Avoid division by zero
        batch_embeddings = sum_embeddings / sum_mask
        return batch_embeddings
d_embeddings = model.model.embed_tokens.weight.shape[1]
embeddings = torch.zeros([0, d_embeddings]).to(device)
for batch in enumerate_batch(train, batch_size):
    stimulii = [x[0] for x in batch]
    tokenized = tokenizer(
        stimulii,
        return_tensors="pt",
        padding=True)
    tokenized.to(device)
    batch_embeddings = embed_batch(**tokenized)
    embeddings = torch.cat([embeddings, batch_embeddings], dim=0)
    print(f"[embedding] {str(embeddings.shape)}")
import hnswlib
embeddings = F.normalize(embeddings, p=2, dim=1)
n_embeddings = embeddings.shape[0]
index_vectors = embeddings.cpu().numpy()
index_indices = [x for x in range(n_embeddings)]
index = hnswlib.Index(space='l2', dim=d_embeddings)
index.init_index(max_elements=n_embeddings, ef_construction=100, M=16)
index.add_items(index_vectors, index_indices)
index.set_ef(50)
print(f"[embedding] complete")
test = load_csv('yelp_review_polarity_csv/test.csv')[:n_test]
instruction = 'Classify the given text as either negative (`1`) or positive (`2`). Write nothing else.'
points = 0
eot_id = tokenizer.convert_tokens_to_ids("<|eot_id|>")
progress = 0
for batch in enumerate_batch(test, batch_size):
    stimulii = [x[0] for x in batch]
    tokenized = tokenizer(
        stimulii,
        return_tensors='pt',
        padding=True)
    tokenized.to(device)
    batch_embeddings = embed_batch(**tokenized).cpu()
    prompts: list[str] = []
    for i in range(len(stimulii)):
        query = batch_embeddings[i,]
        stimulus = stimulii[i]
        labels, distances = index.knn_query(
            query,
            k=n_examples)
        l = labels.squeeze(0).tolist()
        ex_keys = [train[i][0] for i in l]
        ex_values = [train[i][1] for i in l]
        ex_renders = [f"# Example {i + 1}\n\n```\n{ex_keys[i]}\n```\n\n{ex_values[i]}" for i in range(len(ex_keys))]
        ex_render = '\n\n'.join(ex_renders)
        messages = [
            {
                "role": "system",
                "content": f"{instruction}\n\n{ex_render}"
            },
            {
                "role": "user",
                "content": stimulus
            }]
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        prompts.append(prompt)
    prompts_tokenized = tokenizer(
        prompts,
        return_tensors='pt',
        padding=True).to(device)
    with torch.no_grad():
        outputs = model(**prompts_tokenized)
    input_ids = prompts_tokenized['input_ids']
    attention_mask = prompts_tokenized['attention_mask']
    lengths = attention_mask.sum(dim=1)
    last_token_indices = lengths - 1
    batch_indices = torch.arange(input_ids.size(0), device=input_ids.device)
    logits   = outputs.logits[batch_indices, last_token_indices, :].cpu()
    values, indices = logits.max(dim=-1)
    targets = [x[1] for x in batch]
    tokenized_targets = tokenizer(targets, return_tensors="pt", add_special_tokens=False)['input_ids'].squeeze(1).tolist()
    matches = [x[0] == x[1] for x in zip(indices, tokenized_targets)]
    count = len([x for x in matches if x])
    print(f"{indices} vs. {tokenized_targets}; {count}; {progress} / {len(test)}")
    points = points + count
    progress += len(batch)
print(f"[total] {points} / {len(test)}")