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 of the train set to gather into the table?
n_test = 38000 # how much of the test set to run?
batch_size = 16 # adjust 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/minimal.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)
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]
points = 0
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()
    indices: list[int] = []
    for i in range(len(stimulii)):
        query = batch_embeddings[i,]
        stimulus = stimulii[i]
        labels, distances = index.knn_query(
            query,
            k=1)
        l = labels.squeeze(0).tolist()
        i = l[0]
        value = train[i]
    targets = [x[1] for x in batch]
    matches = [x[0] == x[1] for x in zip(indices, 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)}")

<script src="https://gist.github.com/JoeCooper/0ab8fb7eb6613df3646f789e28ddc00b.js"></script>

****<script src="https://gist.github.com/JoeCooper/0ab8fb7eb6613df3646f789e28ddc00b.js"></script>

<script src="https://gist.github.com/JoeCooper/0ab8fb7eb6613df3646f789e28ddc00b.js"></script>

<script

_ @_