<jupyter_start><jupyter_text>Create dummy dataset<jupyter_code>X, Y = make_classification(n_samples=100_000, n_features=256, n_informative=256, n_redundant=0, 
                           n_classes=10_000)

X = th.from_numpy(X).float()
Y = th.from_numpy(Y).long()<jupyter_output><empty_output><jupyter_text>Create the model<jupyter_code>def index_factory(d):
    index = faiss.index_factory(d, "IVF256,Flat", faiss.METRIC_INNER_PRODUCT)
    index.nprobe = 32

    return index

model = th.nn.Sequential(
    th.nn.Linear(256, 256),
    th.nn.ReLU(),
    ApproximateLinear(256, 10_000, index_factory=index_factory),
)

dset   = TensorDataset(X, Y)
loader = DataLoader(dset, batch_size=64)
adam   = th.optim.Adam(model.parameters())<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>for x, y in tqdm(loader):
    adam.zero_grad()
    
    x, y = V(x), V(y)
    o    = model(x)
    
    loss = th.nn.functional.nll_loss(o, y)
    loss.backward()
    
    adam.step()<jupyter_output><empty_output><jupyter_text>Evaluate full<jupyter_code>preds, targets = [], []
loader = DataLoader(dset, batch_size=64, shuffle=False)

with th.autograd.no_grad():
    for x, y in tqdm(loader):
        x, y = V(x), V(y)
        o    = model(x)
        
        p = o.max(1)[1]
        
        preds.append(p)
        targets.append(y)
        
preds   = th.cat(preds)
targets = th.cat(targets)
acc = float((preds == targets).float().sum() / float(preds.size(0)))
print(f'{acc:.5f}')<jupyter_output>0.00012<jupyter_text>Evaluate approximate<jupyter_code>preds, targets = [], []
loader = DataLoader(dset, batch_size=64, shuffle=False)
model  = model.eval()

with th.autograd.no_grad():
    for x, y in tqdm(loader):
        x, y = V(x), V(y)
        _, p = model(x)
        
        preds.append(p.squeeze())
        targets.append(y.squeeze())
        
preds   = th.cat(preds)
targets = th.cat(targets)
acc = float((preds == targets).float().sum() / float(preds.size(0)))
print(f'{acc:.5f}')<jupyter_output>0.00013
{'fraction_non_alphanumeric': 0.15768463073852296, 'fraction_numerical': 0.030439121756487025, 'mean_word_length': 3.50561797752809, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '124125', 'license': 'MIT', 'n_tokens_mistral': 855, 'n_tokens_neox': 798, 'n_words': 120, 'path': 'python/examples/plugins-torch.ipynb', 'repo_name': 'elanmart/mips', 'size': '2004'}