fizban99
commited on
Commit
•
a46f28b
1
Parent(s):
cdee2f2
wikipedia added
Browse files- app.py +13 -3
- mystore/documents.h5 +3 -0
- mystore/embeddings.h5 +3 -0
- mystore/metadatas.h5 +3 -0
- requirements.txt +2 -1
- simiandb.py +339 -0
app.py
CHANGED
@@ -5,9 +5,19 @@ Created on Wed Mar 22 19:59:54 2023
|
|
5 |
"""
|
6 |
|
7 |
import gradio as gr
|
|
|
|
|
8 |
|
9 |
-
def greet(name):
|
10 |
-
return "Hello " + name + "!!"
|
11 |
|
12 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
13 |
iface.launch()
|
|
|
5 |
"""
|
6 |
|
7 |
import gradio as gr
|
8 |
+
from simiandb import Simiandb
|
9 |
+
from langchain.embeddings import HuggingFaceEmbeddings
|
10 |
|
|
|
|
|
11 |
|
12 |
+
|
13 |
+
|
14 |
+
model_name = "all-MiniLM-L6-v2"
|
15 |
+
hf = HuggingFaceEmbeddings(model_name=model_name)
|
16 |
+
|
17 |
+
documentdb = Simiandb("mystore", embedding_function=hf, mode="a")
|
18 |
+
|
19 |
+
def search(query):
|
20 |
+
return documentdb.similarity_search(query)
|
21 |
+
|
22 |
+
iface = gr.Interface(fn=search, inputs="text", outputs="text")
|
23 |
iface.launch()
|
mystore/documents.h5
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:3fc910df8f981cd0d8b7f006a57a09181911d9a6193cda39f82839878bb2f5bb
|
3 |
+
size 1192913753
|
mystore/embeddings.h5
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:e96f3ebb7d86ecc9a4a27207e06115ca542422da9547839ffc403cb70653a71a
|
3 |
+
size 2480735000
|
mystore/metadatas.h5
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:7f233743a5a99181675ebc45f36248e52d36a12238f0f9330869a9c4e967fb0d
|
3 |
+
size 1024
|
requirements.txt
CHANGED
@@ -1,2 +1,3 @@
|
|
1 |
-
|
|
|
2 |
|
|
|
1 |
+
tables
|
2 |
+
langchain
|
3 |
|
simiandb.py
ADDED
@@ -0,0 +1,339 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# -*- coding: utf-8 -*-
|
2 |
+
import tables
|
3 |
+
from pathlib import Path
|
4 |
+
import numpy as np
|
5 |
+
from numpy.lib.recfunctions import structured_to_unstructured
|
6 |
+
from tqdm import tqdm
|
7 |
+
from numba import njit, prange
|
8 |
+
from time import time
|
9 |
+
|
10 |
+
|
11 |
+
|
12 |
+
@njit('float32[:](uint8[:])', parallel=True)
|
13 |
+
def tofp32n8(arr):
|
14 |
+
"""Numba-optimized function that converts a fp8 (4M3E) array to fp32 using a mapping table
|
15 |
+
The array is assumed to be one dimensional with the fp8
|
16 |
+
represented as UInt8
|
17 |
+
"""
|
18 |
+
fp8table= np.frombuffer(b'\x00\x00\x00\x00\x00\x00\x00;\x00\x00\x80;\x00\x00\xc0;\x00\x00\x00<\x00\x00 <\x00\x00@<\x00\x00`<\x00\x00\x80<\x00\x00\x90<\x00\x00\xa0<\x00\x00\xb0<\x00\x00\xc0<\x00\x00\xd0<\x00\x00\xe0<\x00\x00\xf0<\x00\x00\x00=\x00\x00\x10=\x00\x00 =\x00\x000=\x00\x00@=\x00\x00P=\x00\x00`=\x00\x00p=\x00\x00\x80=\x00\x00\x90=\x00\x00\xa0=\x00\x00\xb0=\x00\x00\xc0=\x00\x00\xd0=\x00\x00\xe0=\x00\x00\xf0=\x00\x00\x00>\x00\x00\x10>\x00\x00 >\x00\x000>\x00\x00@>\x00\x00P>\x00\x00`>\x00\x00p>\x00\x00\x80>\x00\x00\x90>\x00\x00\xa0>\x00\x00\xb0>\x00\x00\xc0>\x00\x00\xd0>\x00\x00\xe0>\x00\x00\xf0>\x00\x00\x00?\x00\x00\x10?\x00\x00 ?\x00\x000?\x00\x00@?\x00\x00P?\x00\x00`?\x00\x00p?\x00\x00\x80?\x00\x00\x90?\x00\x00\xa0?\x00\x00\xb0?\x00\x00\xc0?\x00\x00\xd0?\x00\x00\xe0?\x00\x00\xf0?\x00\x00\x00@\x00\x00\x10@\x00\x00 @\x00\x000@\x00\x00@@\x00\x00P@\x00\x00`@\x00\x00p@\x00\x00\x80@\x00\x00\x90@\x00\x00\xa0@\x00\x00\xb0@\x00\x00\xc0@\x00\x00\xd0@\x00\x00\xe0@\x00\x00\xf0@\x00\x00\x00A\x00\x00\x10A\x00\x00 A\x00\x000A\x00\x00@A\x00\x00PA\x00\x00`A\x00\x00pA\x00\x00\x80A\x00\x00\x90A\x00\x00\xa0A\x00\x00\xb0A\x00\x00\xc0A\x00\x00\xd0A\x00\x00\xe0A\x00\x00\xf0A\x00\x00\x00B\x00\x00\x10B\x00\x00 B\x00\x000B\x00\x00@B\x00\x00PB\x00\x00`B\x00\x00pB\x00\x00\x80B\x00\x00\x90B\x00\x00\xa0B\x00\x00\xb0B\x00\x00\xc0B\x00\x00\xd0B\x00\x00\xe0B\x00\x00\xf0B\x00\x00\x00C\x00\x00\x10C\x00\x00 C\x00\x000C\x00\x00@C\x00\x00PC\x00\x00`C\x00\x00pC\x00\x00\x80C\x00\x00\x90C\x00\x00\xa0C\x00\x00\xb0C\x00\x00\xc0C\x00\x00\xd0C\x00\x00\xe0C\x00\x00\xf0C\x00\x00\x00\x80\x00\x00\x00\xbb\x00\x00\x80\xbb\x00\x00\xc0\xbb\x00\x00\x00\xbc\x00\x00 \xbc\x00\x00@\xbc\x00\x00`\xbc\x00\x00\x80\xbc\x00\x00\x90\xbc\x00\x00\xa0\xbc\x00\x00\xb0\xbc\x00\x00\xc0\xbc\x00\x00\xd0\xbc\x00\x00\xe0\xbc\x00\x00\xf0\xbc\x00\x00\x00\xbd\x00\x00\x10\xbd\x00\x00 \xbd\x00\x000\xbd\x00\x00@\xbd\x00\x00P\xbd\x00\x00`\xbd\x00\x00p\xbd\x00\x00\x80\xbd\x00\x00\x90\xbd\x00\x00\xa0\xbd\x00\x00\xb0\xbd\x00\x00\xc0\xbd\x00\x00\xd0\xbd\x00\x00\xe0\xbd\x00\x00\xf0\xbd\x00\x00\x00\xbe\x00\x00\x10\xbe\x00\x00 \xbe\x00\x000\xbe\x00\x00@\xbe\x00\x00P\xbe\x00\x00`\xbe\x00\x00p\xbe\x00\x00\x80\xbe\x00\x00\x90\xbe\x00\x00\xa0\xbe\x00\x00\xb0\xbe\x00\x00\xc0\xbe\x00\x00\xd0\xbe\x00\x00\xe0\xbe\x00\x00\xf0\xbe\x00\x00\x00\xbf\x00\x00\x10\xbf\x00\x00 \xbf\x00\x000\xbf\x00\x00@\xbf\x00\x00P\xbf\x00\x00`\xbf\x00\x00p\xbf\x00\x00\x80\xbf\x00\x00\x90\xbf\x00\x00\xa0\xbf\x00\x00\xb0\xbf\x00\x00\xc0\xbf\x00\x00\xd0\xbf\x00\x00\xe0\xbf\x00\x00\xf0\xbf\x00\x00\x00\xc0\x00\x00\x10\xc0\x00\x00 \xc0\x00\x000\xc0\x00\x00@\xc0\x00\x00P\xc0\x00\x00`\xc0\x00\x00p\xc0\x00\x00\x80\xc0\x00\x00\x90\xc0\x00\x00\xa0\xc0\x00\x00\xb0\xc0\x00\x00\xc0\xc0\x00\x00\xd0\xc0\x00\x00\xe0\xc0\x00\x00\xf0\xc0\x00\x00\x00\xc1\x00\x00\x10\xc1\x00\x00 \xc1\x00\x000\xc1\x00\x00@\xc1\x00\x00P\xc1\x00\x00`\xc1\x00\x00p\xc1\x00\x00\x80\xc1\x00\x00\x90\xc1\x00\x00\xa0\xc1\x00\x00\xb0\xc1\x00\x00\xc0\xc1\x00\x00\xd0\xc1\x00\x00\xe0\xc1\x00\x00\xf0\xc1\x00\x00\x00\xc2\x00\x00\x10\xc2\x00\x00 \xc2\x00\x000\xc2\x00\x00@\xc2\x00\x00P\xc2\x00\x00`\xc2\x00\x00p\xc2\x00\x00\x80\xc2\x00\x00\x90\xc2\x00\x00\xa0\xc2\x00\x00\xb0\xc2\x00\x00\xc0\xc2\x00\x00\xd0\xc2\x00\x00\xe0\xc2\x00\x00\xf0\xc2\x00\x00\x00\xc3\x00\x00\x10\xc3\x00\x00 \xc3\x00\x000\xc3\x00\x00@\xc3\x00\x00P\xc3\x00\x00`\xc3\x00\x00p\xc3\x00\x00\x80\xc3\x00\x00\x90\xc3\x00\x00\xa0\xc3\x00\x00\xb0\xc3\x00\x00\xc0\xc3\x00\x00\xd0\xc3\x00\x00\xe0\xc3\x00\x00\x00\x00\x00\x00\x00;\x00\x00\x80;\x00\x00\xc0;\x00\x00\x00<\x00\x00 <\x00\x00@<\x00\x00`<\x00\x00\x80<\x00\x00\x90<\x00\x00\xa0<\x00\x00\xb0<\x00\x00\xc0<\x00\x00\xd0<\x00\x00\xe0<\x00\x00\xf0<\x00\x00\x00=\x00\x00\x10=\x00\x00 =\x00\x000=\x00\x00@=\x00\x00P=\x00\x00`=\x00\x00p=\x00\x00\x80=\x00\x00\x90=\x00\x00\xa0=\x00\x00\xb0=\x00\x00\xc0=\x00\x00\xd0=\x00\x00\xe0=\x00\x00\xf0=\x00\x00\x00>\x00\x00\x10>\x00\x00 >\x00\x000>\x00\x00@>\x00\x00P>\x00\x00`>\x00\x00p>\x00\x00\x80>\x00\x00\x90>\x00\x00\xa0>\x00\x00\xb0>\x00\x00\xc0>\x00\x00\xd0>\x00\x00\xe0>\x00\x00\xf0>\x00\x00\x00?\x00\x00\x10?\x00\x00 ?\x00\x000?\x00\x00@?\x00\x00P?\x00\x00`?\x00\x00p?\x00\x00\x80?\x00\x00\x90?\x00\x00\xa0?\x00\x00\xb0?\x00\x00\xc0?\x00\x00\xd0?\x00\x00\xe0?\x00\x00\xf0?\x00\x00\x00@\x00\x00\x10@\x00\x00 @\x00\x000@\x00\x00@@\x00\x00P@\x00\x00`@\x00\x00p@\x00\x00\x80@\x00\x00\x90@\x00\x00\xa0@\x00\x00\xb0@\x00\x00\xc0@\x00\x00\xd0@\x00\x00\xe0@\x00\x00\xf0@\x00\x00\x00A\x00\x00\x10A\x00\x00 A\x00\x000A\x00\x00@A\x00\x00PA\x00\x00`A\x00\x00pA\x00\x00\x80A\x00\x00\x90A\x00\x00\xa0A\x00\x00\xb0A\x00\x00\xc0A\x00\x00\xd0A\x00\x00\xe0A\x00\x00\xf0A\x00\x00\x00B\x00\x00\x10B\x00\x00 B\x00\x000B\x00\x00@B\x00\x00PB\x00\x00`B\x00\x00pB\x00\x00\x80B\x00\x00\x90B\x00\x00\xa0B\x00\x00\xb0B\x00\x00\xc0B\x00\x00\xd0B\x00\x00\xe0B\x00\x00\xf0B\x00\x00\x00C\x00\x00\x10C\x00\x00 C\x00\x000C\x00\x00@C\x00\x00PC\x00\x00`C\x00\x00pC\x00\x00\x80C\x00\x00\x90C\x00\x00\xa0C\x00\x00\xb0C\x00\x00\xc0C\x00\x00\xd0C\x00\x00\xe0C\x00\x00\xf0C\x00\x00\x00\x80\x00\x00\x00\xbb\x00\x00\x80\xbb\x00\x00\xc0\xbb\x00\x00\x00\xbc\x00\x00 \xbc\x00\x00@\xbc\x00\x00`\xbc\x00\x00\x80\xbc\x00\x00\x90\xbc\x00\x00\xa0\xbc\x00\x00\xb0\xbc\x00\x00\xc0\xbc\x00\x00\xd0\xbc\x00\x00\xe0\xbc\x00\x00\xf0\xbc\x00\x00\x00\xbd\x00\x00\x10\xbd\x00\x00 \xbd\x00\x000\xbd\x00\x00@\xbd\x00\x00P\xbd\x00\x00`\xbd\x00\x00p\xbd\x00\x00\x80\xbd\x00\x00\x90\xbd\x00\x00\xa0\xbd\x00\x00\xb0\xbd\x00\x00\xc0\xbd\x00\x00\xd0\xbd\x00\x00\xe0\xbd\x00\x00\xf0\xbd\x00\x00\x00\xbe\x00\x00\x10\xbe\x00\x00 \xbe\x00\x000\xbe\x00\x00@\xbe\x00\x00P\xbe\x00\x00`\xbe\x00\x00p\xbe\x00\x00\x80\xbe\x00\x00\x90\xbe\x00\x00\xa0\xbe\x00\x00\xb0\xbe\x00\x00\xc0\xbe\x00\x00\xd0\xbe\x00\x00\xe0\xbe\x00\x00\xf0\xbe\x00\x00\x00\xbf\x00\x00\x10\xbf\x00\x00 \xbf\x00\x000\xbf\x00\x00@\xbf\x00\x00P\xbf\x00\x00`\xbf\x00\x00p\xbf\x00\x00\x80\xbf\x00\x00\x90\xbf\x00\x00\xa0\xbf\x00\x00\xb0\xbf\x00\x00\xc0\xbf\x00\x00\xd0\xbf\x00\x00\xe0\xbf\x00\x00\xf0\xbf\x00\x00\x00\xc0\x00\x00\x10\xc0\x00\x00 \xc0\x00\x000\xc0\x00\x00@\xc0\x00\x00P\xc0\x00\x00`\xc0\x00\x00p\xc0\x00\x00\x80\xc0\x00\x00\x90\xc0\x00\x00\xa0\xc0\x00\x00\xb0\xc0\x00\x00\xc0\xc0\x00\x00\xd0\xc0\x00\x00\xe0\xc0\x00\x00\xf0\xc0\x00\x00\x00\xc1\x00\x00\x10\xc1\x00\x00 \xc1\x00\x000\xc1\x00\x00@\xc1\x00\x00P\xc1\x00\x00`\xc1\x00\x00p\xc1\x00\x00\x80\xc1\x00\x00\x90\xc1\x00\x00\xa0\xc1\x00\x00\xb0\xc1\x00\x00\xc0\xc1\x00\x00\xd0\xc1\x00\x00\xe0\xc1\x00\x00\xf0\xc1\x00\x00\x00\xc2\x00\x00\x10\xc2\x00\x00 \xc2\x00\x000\xc2\x00\x00@\xc2\x00\x00P\xc2\x00\x00`\xc2\x00\x00p\xc2\x00\x00\x80\xc2\x00\x00\x90\xc2\x00\x00\xa0\xc2\x00\x00\xb0\xc2\x00\x00\xc0\xc2\x00\x00\xd0\xc2\x00\x00\xe0\xc2\x00\x00\xf0\xc2\x00\x00\x00\xc3\x00\x00\x10\xc3\x00\x00 \xc3\x00\x000\xc3\x00\x00@\xc3\x00\x00P\xc3\x00\x00`\xc3\x00\x00p\xc3\x00\x00\x80\xc3\x00\x00\x90\xc3\x00\x00\xa0\xc3\x00\x00\xb0\xc3\x00\x00\xc0\xc3\x00\x00\xd0\xc3\x00\x00\xe0\xc3', dtype=np.float32)
|
19 |
+
arr2 = np.empty(arr.shape[0], dtype="float32")
|
20 |
+
for i in prange(arr.shape[0]):
|
21 |
+
arr2[i] = fp8table[arr[i]]
|
22 |
+
return arr2
|
23 |
+
|
24 |
+
|
25 |
+
def tofp32(arr):
|
26 |
+
"""Converts a fp8 (4M3E) array to fp32.
|
27 |
+
Reshapes the array to be one
|
28 |
+
dimensional and uses a numba-optimized function
|
29 |
+
"""
|
30 |
+
return tofp32n8(arr.reshape(arr.shape[0]*arr.shape[1])).reshape(arr.shape)
|
31 |
+
|
32 |
+
|
33 |
+
@njit('uint8[:](uint32[:])', parallel=True)
|
34 |
+
def tofp8n(arr):
|
35 |
+
"""Numba-optimized function that converts an array of fp32 to fp8 (4M3E)
|
36 |
+
Uses the algorithm described by ProjectPhysX at https://stackoverflow.com/questions/1659440/32-bit-to-16-bit-floating-point-conversion
|
37 |
+
and https://www.researchgate.net/publication/362275548_Accuracy_and_performance_of_the_lattice_Boltzmann_method_with_64-bit_32-bit_and_customized_16-bit_number_formats
|
38 |
+
"""
|
39 |
+
arr2 = np.empty(arr.shape[0], dtype="uint8")
|
40 |
+
for i in prange(arr.shape[0]):
|
41 |
+
# round-to-nearest-even: add last bit after truncated mantissa (1+8+3) from left
|
42 |
+
y = arr[i] + 0x00080000
|
43 |
+
e = (y&0x7F800000)>>23 # exponent
|
44 |
+
m = y&0x007FFFFF #mantissa
|
45 |
+
|
46 |
+
if e > 135:
|
47 |
+
arr2[i] = 0x7F | (y&0x80000000)>>24 # saturated
|
48 |
+
elif e > 120:
|
49 |
+
arr2[i] = ((e-120)<<3) & 0x78 | m>>20 | (y&0x80000000)>>24 # normalized
|
50 |
+
elif e < 121 and e > 116:
|
51 |
+
# 0x00780000 = 0x00800000-0x00080000 = decimal indicator flag - initial rounding
|
52 |
+
arr2[i] = ((((m+0x00780000)>>(140-e))+1)>>1) | (y&0x80000000)>>24
|
53 |
+
else:
|
54 |
+
arr2[i] = 0 | (y&0x80000000)>>24
|
55 |
+
return arr2
|
56 |
+
|
57 |
+
|
58 |
+
def tofp8(arr):
|
59 |
+
"""Converts an array of fp32 to fp8 (4M3E)
|
60 |
+
Reshapes the array to be one
|
61 |
+
dimensional and uses a numba-optimized function
|
62 |
+
"""
|
63 |
+
return tofp8n(arr.view(dtype=np.uint32).reshape(arr.shape[0]*arr.shape[1])).view(dtype=np.uint8).reshape(arr.shape)
|
64 |
+
|
65 |
+
|
66 |
+
|
67 |
+
class BlobTable():
|
68 |
+
"""Class to handle a storage of variable-length values of a key-value storage
|
69 |
+
Key is fixed length of key_length
|
70 |
+
"""
|
71 |
+
def __init__(self, store, key_length=20):
|
72 |
+
"""Initializes class using a pytables store and a key_length value
|
73 |
+
"""
|
74 |
+
if "keys" not in store.root:
|
75 |
+
# reasonable compression optimized for reading speed
|
76 |
+
filters = tables.Filters(complevel=5, complib='blosc:lz4',
|
77 |
+
shuffle=1, bitshuffle=0)
|
78 |
+
|
79 |
+
blob_type = {"key": tables.StringCol(key_length, pos=0),
|
80 |
+
"offset":tables.Int64Col(pos=1),
|
81 |
+
"length": tables.Int64Col(pos=2),
|
82 |
+
}
|
83 |
+
|
84 |
+
self.keys_table = store.create_table("/", "keys",
|
85 |
+
blob_type,
|
86 |
+
filters=filters,
|
87 |
+
chunkshape=10000)
|
88 |
+
self.values_table = store.create_earray("/", "values", atom=tables.UInt8Atom(), shape=(0,), filters=filters)
|
89 |
+
else:
|
90 |
+
self.keys_table = store.root.keys
|
91 |
+
self.values_table = store.root.values
|
92 |
+
|
93 |
+
self.offset = self.values_table.nrows
|
94 |
+
self.nrows = self.keys_table.nrows
|
95 |
+
self._is_closed = False
|
96 |
+
|
97 |
+
|
98 |
+
def __len__(self):
|
99 |
+
return self.nrows
|
100 |
+
|
101 |
+
def create_index(self):
|
102 |
+
self.keys_table.cols.key.reindex()
|
103 |
+
|
104 |
+
|
105 |
+
def append(self, key, value):
|
106 |
+
"""Appends a key-value to the storage
|
107 |
+
"""
|
108 |
+
# store variable length value
|
109 |
+
length = len(value)
|
110 |
+
self.values_table.append(np.frombuffer(value, dtype=np.uint8))
|
111 |
+
|
112 |
+
# store index
|
113 |
+
row = self.keys_table.row
|
114 |
+
row["key"] = key
|
115 |
+
row["offset"] = self.offset
|
116 |
+
row["length"] = length
|
117 |
+
row.append()
|
118 |
+
self.offset += length
|
119 |
+
self.nrows += 1
|
120 |
+
|
121 |
+
def __getitem__ (self, rownum):
|
122 |
+
if isinstance(rownum, slice):
|
123 |
+
return [self[ii] for ii in range(*rownum.indices(len(self)))]
|
124 |
+
else:
|
125 |
+
row = self.keys_table[rownum]
|
126 |
+
offset = row['offset']
|
127 |
+
value = self.values_table.read(offset, offset+row["length"]).tobytes()
|
128 |
+
|
129 |
+
return value
|
130 |
+
|
131 |
+
|
132 |
+
def get_value (self, key):
|
133 |
+
key = key.encode("utf8")
|
134 |
+
offset, length = [(r['offset'], r['length']) for r in self.keys_table.where(f"key=={key}")][0]
|
135 |
+
value = self.values_table.read(offset, offset+length).tobytes()
|
136 |
+
return value
|
137 |
+
|
138 |
+
|
139 |
+
class Simiandb():
|
140 |
+
"""Wrapper around pytables store .
|
141 |
+
To use, you should have the ``pytables`` python package installed.
|
142 |
+
Example:
|
143 |
+
.. code-block:: python
|
144 |
+
from simiandb import Simiandb
|
145 |
+
docdb = simiandb("store")
|
146 |
+
"""
|
147 |
+
|
148 |
+
def __init__(self, storepath, embedding_function=None, mode="a", id_length = 19):
|
149 |
+
|
150 |
+
if mode not in ["a", "w", "r"]:
|
151 |
+
raise ValueError("Mode can only be r, w or a")
|
152 |
+
self._embedding_function = embedding_function
|
153 |
+
self._storename = Path(storepath)
|
154 |
+
self._mode = mode
|
155 |
+
if not self._storename.exists():
|
156 |
+
self._storename.mkdir()
|
157 |
+
|
158 |
+
self._vectorstore = tables.open_file( self._storename / "embeddings.h5", mode = mode)
|
159 |
+
self._docstore = tables.open_file( self._storename / "documents.h5", mode = mode)
|
160 |
+
self._metastore = tables.open_file( self._storename / "metadatas.h5", mode = mode)
|
161 |
+
self._embedding_function = embedding_function
|
162 |
+
self._is_closed = False
|
163 |
+
if 'embeddings' in self._vectorstore.root:
|
164 |
+
self._vector_table = self._vectorstore.root.embeddings
|
165 |
+
self._docs_table = BlobTable(self._docstore, id_length)
|
166 |
+
return
|
167 |
+
|
168 |
+
|
169 |
+
def __enter__(self):
|
170 |
+
"""Magic method Required for usage with the with statement
|
171 |
+
"""
|
172 |
+
return self
|
173 |
+
|
174 |
+
|
175 |
+
def _get_top_indexes(self, c, k):
|
176 |
+
count = self._vector_table.nrows
|
177 |
+
st =0
|
178 |
+
batch = self._vector_table.chunkshape[0]*25
|
179 |
+
res = np.ascontiguousarray(np.empty(shape=(count,), dtype="float32"))
|
180 |
+
end = 0
|
181 |
+
a = time()
|
182 |
+
while end!=count:
|
183 |
+
end += batch
|
184 |
+
end = end if end <= count else count
|
185 |
+
t_res = structured_to_unstructured(self._vector_table.read(start=st, stop=end))
|
186 |
+
t_res = tofp32(t_res)
|
187 |
+
np.dot(t_res,c, res[st:end])
|
188 |
+
st = end
|
189 |
+
|
190 |
+
indices = np.argpartition(res, -k)[-k:] #from https://stackoverflow.com/questions/6910641/how-do-i-get-indices-of-n-maximum-values-in-a-numpy-array
|
191 |
+
indices = indices[np.argsort(res[indices])[::-1]]
|
192 |
+
print(time() -a)
|
193 |
+
return indices
|
194 |
+
|
195 |
+
|
196 |
+
def _create_embeddings_table(self, dimensions):
|
197 |
+
"""Creates the embeddings table within the pytables file
|
198 |
+
"""
|
199 |
+
if dimensions > 512:
|
200 |
+
# prevent pytables warning on max_columns
|
201 |
+
tables.parameters.MAX_COLUMNS = len(dimensions)
|
202 |
+
embedding_type = {f"d{n}":tables.UInt8Col(pos=n) for n in range(dimensions)}
|
203 |
+
|
204 |
+
# no compression for embeddings
|
205 |
+
filters = None
|
206 |
+
|
207 |
+
self._vector_table = self._vectorstore.create_table("/", "embeddings",
|
208 |
+
embedding_type,
|
209 |
+
filters=filters,
|
210 |
+
chunkshape=10000)
|
211 |
+
|
212 |
+
|
213 |
+
def _check_closed(self):
|
214 |
+
if self._is_closed:
|
215 |
+
raise ValueError("Simiandb is already closed")
|
216 |
+
|
217 |
+
|
218 |
+
def add_texts(self, texts, metadatas = None, ids = None, embeddings=None, show_progressbar=True):
|
219 |
+
"""Run more texts through the embeddings and add to the vectorstore.
|
220 |
+
Args:
|
221 |
+
texts (Iterable[str]): Texts to add to the vectorstore.
|
222 |
+
metadatas (Optional[List[dict]], optional): Optional list of metadatas.
|
223 |
+
ids (Optional[List[str]], optional): Optional list of IDs.
|
224 |
+
embeddings (Optional[List[array]], optional): Optional list of embeddings.
|
225 |
+
Returns:
|
226 |
+
List[str]: List of IDs of the added texts.
|
227 |
+
"""
|
228 |
+
|
229 |
+
self._check_closed()
|
230 |
+
|
231 |
+
self._add_embeddings(texts, embeddings, show_progressbar)
|
232 |
+
|
233 |
+
if ids is None:
|
234 |
+
ids = list(range(self.docs_table.nrows, self.docs_table.nrows + len(texts)))
|
235 |
+
|
236 |
+
for textid, text in zip(ids, texts):
|
237 |
+
self.docs_table.append(textid, text.encode("utf8"))
|
238 |
+
|
239 |
+
return ids
|
240 |
+
|
241 |
+
|
242 |
+
def get_text(self, key):
|
243 |
+
return self._docs_table.get_value(key).decode("utf8")
|
244 |
+
|
245 |
+
|
246 |
+
def create_keys_index(self):
|
247 |
+
self._docs_table.create_index()
|
248 |
+
|
249 |
+
|
250 |
+
def _add_embeddings(self, texts, embeddings, show_progressbar):
|
251 |
+
"""Calculate or use embeddings to fill the embeddings table
|
252 |
+
"""
|
253 |
+
|
254 |
+
if embeddings is None and not self._embedding_function is None:
|
255 |
+
embeddings = self._embedding_function.embed_documents(texts)
|
256 |
+
|
257 |
+
if not embeddings is None and 'embeddings' not in self._vectorstore.root:
|
258 |
+
dimensions = len(embeddings[0])
|
259 |
+
self._create_embeddings_table(dimensions)
|
260 |
+
|
261 |
+
if not embeddings is None :
|
262 |
+
self._vector_table = self._vectorstore.root.embeddings
|
263 |
+
embeddings = tofp8(np.array(embeddings, dtype=np.float32))
|
264 |
+
self._vector_table.append(embeddings)
|
265 |
+
|
266 |
+
|
267 |
+
def regenerate_embeddings(self, embeddings=None, show_progressbar=True):
|
268 |
+
"""Run existing texts through the embeddings and add to the vectorstore.
|
269 |
+
Args:
|
270 |
+
embeddings (Optional[List[array]], optional): Optional list of embeddings.
|
271 |
+
"""
|
272 |
+
|
273 |
+
self._check_closed()
|
274 |
+
self._vectorstore.close()
|
275 |
+
(self._storename / "embeddings.h5").kill()
|
276 |
+
self._vectorstore = tables.open_file( self._storename / "embeddings.h5", mode = self._mode)
|
277 |
+
|
278 |
+
batch_size = 1000
|
279 |
+
for i in tqdm(range(0, len(self.docs_table), batch_size), disable=not show_progressbar):
|
280 |
+
text_batch = [text.decode("utf8") for text in self.docs_table[i:i+batch_size]]
|
281 |
+
if embeddings is not None:
|
282 |
+
embeddings_batch = embeddings[i:i+batch_size]
|
283 |
+
elif self.embedding_function is not None:
|
284 |
+
embeddings_batch = self._embedding_function.embed_documents(text_batch)
|
285 |
+
else:
|
286 |
+
raise ValueError("Neither embeddings nor embedding function provided")
|
287 |
+
self._add_embeddings(text_batch, embeddings_batch, show_progressbar)
|
288 |
+
return
|
289 |
+
|
290 |
+
|
291 |
+
def similarity_search(self, query: str, k = 4, filter = None):
|
292 |
+
"""Run similarity search with PytableStore.
|
293 |
+
Args:
|
294 |
+
query (str): Query text to search for.
|
295 |
+
k (int): Number of results to return. Defaults to 4.
|
296 |
+
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
|
297 |
+
Returns:
|
298 |
+
List[Document]: List of documents most simmilar to the query text.
|
299 |
+
"""
|
300 |
+
self._check_closed()
|
301 |
+
query_embedding = np.array(self._embedding_function.embed_query(query),dtype="float32")
|
302 |
+
results = self._get_top_indexes(query_embedding, k)
|
303 |
+
|
304 |
+
docs = [self._docs_table[i].decode("utf8") for i in results]
|
305 |
+
return docs
|
306 |
+
|
307 |
+
|
308 |
+
def close(self):
|
309 |
+
"""Makes sure the pytables file is closed
|
310 |
+
"""
|
311 |
+
if not self._is_closed:
|
312 |
+
self._is_closed = True
|
313 |
+
|
314 |
+
if hasattr(self, '_Simiandb__vectorstore'):
|
315 |
+
try:
|
316 |
+
self._vectorstore.flush()
|
317 |
+
self._docstore.flush()
|
318 |
+
self._metastore.flush()
|
319 |
+
self._vectorstore.close()
|
320 |
+
self._docstore.close()
|
321 |
+
self._metastore.close()
|
322 |
+
except:
|
323 |
+
print("Unable to close file")
|
324 |
+
|
325 |
+
|
326 |
+
def __exit__(self, exc_type, exc_value, exc_traceback):
|
327 |
+
"""Magic method Required for usage with the with statement
|
328 |
+
"""
|
329 |
+
self.close()
|
330 |
+
|
331 |
+
def __del__(self):
|
332 |
+
"""Magic method just in case the object is deleted without closing it
|
333 |
+
"""
|
334 |
+
self.close()
|
335 |
+
|
336 |
+
|
337 |
+
|
338 |
+
if __name__ == '__main__':
|
339 |
+
pass
|