Spaces:
Runtime error
Runtime error
File size: 2,898 Bytes
1f22b14 71e7dd8 1f22b14 71e7dd8 06bca0c 1f22b14 06bca0c 71e7dd8 1f22b14 71e7dd8 1f22b14 06bca0c 1f22b14 06bca0c 71e7dd8 1f22b14 71e7dd8 1f22b14 71e7dd8 1f22b14 06bca0c 1f22b14 6aad21a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import numpy as np
import pandas as pd
import pytest
from buster.documents import DocumentsDB, DocumentsPickle
from buster.retriever import PickleRetriever, SQLiteRetriever
@pytest.mark.parametrize(
"documents_manager, retriever, extension",
[(DocumentsDB, SQLiteRetriever, "db"), (DocumentsPickle, PickleRetriever, "tar.gz")],
)
def test_write_read(tmp_path, documents_manager, retriever, extension):
db = documents_manager(tmp_path / f"test.{extension}")
data = pd.DataFrame.from_dict(
{
"title": ["test"],
"url": ["http://url.com"],
"content": ["cool text"],
"embedding": [np.arange(10, dtype=np.float32) - 0.3],
"n_tokens": [10],
}
)
db.add(source="test", df=data)
db_data = retriever(tmp_path / f"test.{extension}").get_documents("test")
assert db_data["title"].iloc[0] == data["title"].iloc[0]
assert db_data["url"].iloc[0] == data["url"].iloc[0]
assert db_data["content"].iloc[0] == data["content"].iloc[0]
assert np.allclose(db_data["embedding"].iloc[0], data["embedding"].iloc[0])
assert db_data["n_tokens"].iloc[0] == data["n_tokens"].iloc[0]
@pytest.mark.parametrize(
"documents_manager, retriever, extension",
[(DocumentsDB, SQLiteRetriever, "db"), (DocumentsPickle, PickleRetriever, "tar.gz")],
)
def test_write_write_read(tmp_path, documents_manager, retriever, extension):
db = documents_manager(tmp_path / f"test.{extension}")
data_1 = pd.DataFrame.from_dict(
{
"title": ["test"],
"url": ["http://url.com"],
"content": ["cool text"],
"embedding": [np.arange(10, dtype=np.float32) - 0.3],
"n_tokens": [10],
}
)
db.add(source="test", df=data_1)
data_2 = pd.DataFrame.from_dict(
{
"title": ["other"],
"url": ["http://url.com/page.html"],
"content": ["lorem ipsum"],
"embedding": [np.arange(20, dtype=np.float32) / 10 - 2.3],
"n_tokens": [20],
}
)
db.add(source="test", df=data_2)
db_data = retriever(tmp_path / f"test.{extension}").get_documents("test")
assert len(db_data) == len(data_2)
assert db_data["title"].iloc[0] == data_2["title"].iloc[0]
assert db_data["url"].iloc[0] == data_2["url"].iloc[0]
assert db_data["content"].iloc[0] == data_2["content"].iloc[0]
assert np.allclose(db_data["embedding"].iloc[0], data_2["embedding"].iloc[0])
assert db_data["n_tokens"].iloc[0] == data_2["n_tokens"].iloc[0]
def test_update_source(tmp_path):
display_name = "Super Test"
db = DocumentsDB(tmp_path / "test.db")
db.update_source(source="test", display_name=display_name)
returned_display_name = SQLiteRetriever(tmp_path / "test.db").get_source_display_name("test")
assert display_name == returned_display_name
|