chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,229 @@
|
||||
"""
|
||||
Agent module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from smolagents import CodeAgent, PythonInterpreterTool
|
||||
|
||||
from txtai.agent import Agent
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
# agents.md content
|
||||
AGENTS = """
|
||||
Basic instructions here
|
||||
"""
|
||||
|
||||
# Sample skill.md content
|
||||
SKILL = """---
|
||||
name: hello
|
||||
description: says hello world
|
||||
---
|
||||
|
||||
Says hello world
|
||||
"""
|
||||
|
||||
|
||||
class TestAgent(unittest.TestCase):
|
||||
"""
|
||||
Agent tests.
|
||||
"""
|
||||
|
||||
def testExecute(self):
|
||||
"""
|
||||
Test executing main agent loop
|
||||
"""
|
||||
|
||||
agent = Agent(llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_steps=1)
|
||||
|
||||
# Patch LLM to generate answer
|
||||
agent.process.model.llm = lambda *args, **kwargs: 'Action:\n{"name": "final_answer", "arguments": "Hi"}'
|
||||
|
||||
self.assertEqual(agent("Hello"), "Hi")
|
||||
|
||||
def testInstructions(self):
|
||||
"""
|
||||
Test loading an agents.md file
|
||||
"""
|
||||
|
||||
# Test loading instructions from file
|
||||
agents = os.path.join(tempfile.gettempdir(), "agents.md")
|
||||
with open(agents, "w", encoding="utf-8") as output:
|
||||
output.write(AGENTS)
|
||||
|
||||
agent = Agent(instructions=agents, llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_iterations=1)
|
||||
agent.process.model.llm = lambda *args, **kwargs: 'Action:\n{"name": "final_answer", "arguments": "Hi"}'
|
||||
self.assertEqual(agent("Hello"), "Hi")
|
||||
|
||||
# Test loading from memory
|
||||
agent = Agent(instructions=AGENTS, llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_iterations=1)
|
||||
agent.process.model.llm = lambda *args, **kwargs: 'Action:\n{"name": "final_answer", "arguments": "Hi"}'
|
||||
self.assertEqual(agent("Hello"), "Hi")
|
||||
|
||||
def testMemory(self):
|
||||
"""
|
||||
Test agent memory
|
||||
"""
|
||||
|
||||
agent = Agent(llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_steps=1, memory=5)
|
||||
|
||||
# Patch LLM to generate answer
|
||||
agent.process.model.llm = lambda *args, **kwargs: 'Action:\n{"name": "final_answer", "arguments": "Hi"}'
|
||||
|
||||
self.assertEqual(agent("Hello"), "Hi")
|
||||
self.assertEqual(agent("Hello"), "Hi")
|
||||
|
||||
# Test that results are stored in shared memory
|
||||
self.assertEqual(len(agent.memory.get(None)), 2)
|
||||
|
||||
# Test resetting shared memory
|
||||
self.assertEqual(agent("Hello", reset=True), "Hi")
|
||||
self.assertEqual(len(agent.memory.get(None)), 1)
|
||||
|
||||
# Test session memory
|
||||
self.assertEqual(agent("Hello", session="session-0"), "Hi")
|
||||
self.assertEqual(len(agent.memory.get("session-0")), 1)
|
||||
|
||||
# Test resetting session memory
|
||||
self.assertEqual(agent("Hello", session="session-0", reset=True), "Hi")
|
||||
self.assertEqual(len(agent.memory.get("session-0")), 1)
|
||||
self.assertEqual(len(agent.memory.get(None)), 1)
|
||||
|
||||
def testMethod(self):
|
||||
"""
|
||||
Test agent process methods
|
||||
"""
|
||||
|
||||
agent = Agent(method="code", llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_iterations=1)
|
||||
self.assertIsInstance(agent.process, CodeAgent)
|
||||
|
||||
def testSkill(self):
|
||||
"""
|
||||
Test running a skill from a skill.md file
|
||||
"""
|
||||
|
||||
skill = os.path.join(tempfile.gettempdir(), "skill.md")
|
||||
with open(skill, "w", encoding="utf-8") as output:
|
||||
output.write(SKILL)
|
||||
|
||||
agent = Agent(tools=[skill], llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_iterations=1)
|
||||
|
||||
self.assertIsInstance(agent.tools["hello"]("say hello"), str)
|
||||
|
||||
def testToolsBasic(self):
|
||||
"""
|
||||
Test adding basic function tools
|
||||
"""
|
||||
|
||||
class DateTime:
|
||||
"""
|
||||
Date time instance
|
||||
"""
|
||||
|
||||
def __call__(self, iso):
|
||||
"""
|
||||
Gets the current date and time
|
||||
|
||||
Args:
|
||||
iso: date will be converted to iso format if True
|
||||
|
||||
Returns:
|
||||
current date and time
|
||||
"""
|
||||
|
||||
return datetime.today().isoformat() if iso else datetime.today()
|
||||
|
||||
today = {"name": "today", "description": "Gets the current date and time", "target": DateTime()}
|
||||
|
||||
def current(iso: str) -> str:
|
||||
"""
|
||||
Gets the current date and time
|
||||
|
||||
Args:
|
||||
iso: date will be converted to iso format if True
|
||||
|
||||
Returns:
|
||||
current date and time
|
||||
"""
|
||||
|
||||
return datetime.today().isoformat() if iso else datetime.today()
|
||||
|
||||
agent = Agent(tools=[today, current, "websearch"], llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_steps=1)
|
||||
|
||||
self.assertIsNotNone(agent)
|
||||
self.assertIsInstance(agent.tools["today"](True), str)
|
||||
self.assertIsInstance(agent.tools["current"](True), str)
|
||||
|
||||
def testToolsDefaults(self):
|
||||
"""
|
||||
Test default toolkit tools
|
||||
"""
|
||||
|
||||
agent = Agent(tools=["defaults"], llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_steps=1)
|
||||
|
||||
# Working directory
|
||||
work = tempfile.gettempdir()
|
||||
|
||||
# Test file
|
||||
path = os.path.join(work, "agent_tools.txt")
|
||||
agent.tools["write"](path, "hello world")
|
||||
|
||||
# Test default tools
|
||||
self.assertIsNotNone(agent.tools["bash"](["ls", work]))
|
||||
self.assertGreater(len(agent.tools["glob"](work)), 0)
|
||||
self.assertGreater(len(agent.tools["grep"]("world", "*")), 0)
|
||||
self.assertEqual(agent.tools["todowrite"]("plan"), "plan")
|
||||
|
||||
agent.tools["edit"](path, "hello", "goodbye")
|
||||
self.assertEqual(agent.tools["read"](path), "goodbye world".strip())
|
||||
|
||||
def testToolsEmbeddings(self):
|
||||
"""
|
||||
Test adding Embeddings as a tool
|
||||
"""
|
||||
|
||||
embeddings = Embeddings()
|
||||
embeddings.index(["test"])
|
||||
|
||||
# Generate temp file path and save
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.agent")
|
||||
embeddings.save(index)
|
||||
|
||||
embeddings1 = {
|
||||
"name": "embeddings1",
|
||||
"description": "Searches a test database",
|
||||
"target": embeddings,
|
||||
}
|
||||
|
||||
embeddings2 = {"name": "embeddings2", "description": "Searches a test database", "path": index}
|
||||
|
||||
agent = Agent(tools=[embeddings1, embeddings2], llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_steps=1)
|
||||
|
||||
self.assertIsNotNone(agent)
|
||||
self.assertIsInstance(agent.tools["embeddings1"]("test"), list)
|
||||
|
||||
# pylint: disable=C0115,C0116
|
||||
@patch("mcpadapt.core.MCPAdapt")
|
||||
def testToolsMCP(self, mcp):
|
||||
"""
|
||||
Test adding a MCP tool collection
|
||||
"""
|
||||
|
||||
class MCPAdapt:
|
||||
def __init__(self, *args):
|
||||
self.args = args
|
||||
|
||||
def tools(self):
|
||||
return [PythonInterpreterTool()]
|
||||
|
||||
# Patch MCP adapter for testing
|
||||
mcp.side_effect = MCPAdapt
|
||||
|
||||
agent = Agent(tools=["http://localhost:8000/mcp"], llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_steps=1)
|
||||
self.assertEqual(len(agent.tools), 2)
|
||||
@@ -0,0 +1,537 @@
|
||||
"""
|
||||
Dense ANN module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.ann import ANNFactory, ANN
|
||||
from txtai.serialize import SerializeFactory
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestDense(unittest.TestCase):
|
||||
"""
|
||||
Dense ANN tests.
|
||||
"""
|
||||
|
||||
def testAnnoy(self):
|
||||
"""
|
||||
Test Annoy backend
|
||||
"""
|
||||
|
||||
self.runTests("annoy", None, False)
|
||||
|
||||
def testAnnoyCustom(self):
|
||||
"""
|
||||
Test Annoy backend with custom settings
|
||||
"""
|
||||
|
||||
# Test with custom settings
|
||||
self.runTests("annoy", {"annoy": {"ntrees": 2, "searchk": 1}}, False)
|
||||
|
||||
def testCustomBackend(self):
|
||||
"""
|
||||
Test resolving a custom backend
|
||||
"""
|
||||
|
||||
self.runTests("txtai.ann.Faiss")
|
||||
|
||||
def testCustomBackendNotFound(self):
|
||||
"""
|
||||
Test resolving an unresolvable backend
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ANNFactory.create({"backend": "notfound.ann"})
|
||||
|
||||
def testFaiss(self):
|
||||
"""
|
||||
Test Faiss backend
|
||||
"""
|
||||
|
||||
self.runTests("faiss")
|
||||
|
||||
def testFaissBinary(self):
|
||||
"""
|
||||
Test Faiss backend with a binary hash index
|
||||
"""
|
||||
|
||||
ann = ANNFactory.create({"backend": "faiss", "quantize": 1, "dimensions": 240 * 8, "faiss": {"components": "BHash32"}})
|
||||
|
||||
# Generate and index dummy data
|
||||
data = np.random.rand(100, 240).astype(np.uint8)
|
||||
ann.index(data)
|
||||
|
||||
# Generate query vector and test search
|
||||
query = np.random.rand(240).astype(np.uint8)
|
||||
self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)
|
||||
|
||||
def testFaissCustom(self):
|
||||
"""
|
||||
Test Faiss backend with custom settings
|
||||
"""
|
||||
|
||||
# Test with custom settings
|
||||
self.runTests("faiss", {"faiss": {"nprobe": 2, "components": "PCA16,IDMap,SQ8", "sample": 1.0}}, False)
|
||||
self.runTests("faiss", {"faiss": {"components": "IVF,SQ8"}}, False)
|
||||
|
||||
@patch("platform.system")
|
||||
def testFaissMacOS(self, system):
|
||||
"""
|
||||
Test Faiss backend with macOS
|
||||
"""
|
||||
|
||||
# Run test
|
||||
system.return_value = "Darwin"
|
||||
|
||||
# pylint: disable=C0415, W0611
|
||||
# Force reload of class
|
||||
name = "txtai.ann.dense.faiss"
|
||||
module = sys.modules[name]
|
||||
del sys.modules[name]
|
||||
import txtai.ann.dense.faiss
|
||||
|
||||
# Run tests
|
||||
self.runTests("faiss")
|
||||
|
||||
# Restore original module
|
||||
sys.modules[name] = module
|
||||
|
||||
@unittest.skipIf(os.name == "nt", "mmap not supported on Windows")
|
||||
def testFaissMmap(self):
|
||||
"""
|
||||
Test Faiss backend with mmap enabled
|
||||
"""
|
||||
|
||||
# Test to with mmap enabled
|
||||
self.runTests("faiss", {"faiss": {"mmap": True}}, False)
|
||||
|
||||
def testGGML(self):
|
||||
"""
|
||||
Test GGML backend
|
||||
"""
|
||||
|
||||
self.runTests("ggml")
|
||||
|
||||
def testGGMLQuantization(self):
|
||||
"""
|
||||
Test GGML backend with quantization enabled
|
||||
"""
|
||||
|
||||
ann = ANNFactory.create({"backend": "ggml", "ggml": {"quantize": "Q4_0"}})
|
||||
|
||||
# Generate and index dummy data
|
||||
data = np.random.rand(100, 256).astype(np.float32)
|
||||
ann.index(data)
|
||||
|
||||
# Test save and load
|
||||
index = os.path.join(tempfile.gettempdir(), "ggml.q4_0.v1")
|
||||
ann.save(index)
|
||||
ann.load(index)
|
||||
|
||||
# Generate query vector and test search
|
||||
query = np.random.rand(256).astype(np.float32)
|
||||
self.normalize(query)
|
||||
self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)
|
||||
|
||||
# Validate count
|
||||
self.assertEqual(ann.count(), 100)
|
||||
|
||||
# Test delete
|
||||
ann.delete([0])
|
||||
self.assertEqual(ann.count(), 99)
|
||||
|
||||
# Save updated index with deletes and reload
|
||||
index = os.path.join(tempfile.gettempdir(), "ggml.q4_0.v2")
|
||||
ann.save(index)
|
||||
ann.load(index)
|
||||
ann.index(data)
|
||||
|
||||
def testGGMLInvalid(self):
|
||||
"""
|
||||
Test invalid GGML configurations
|
||||
"""
|
||||
|
||||
data = np.random.rand(100, 240).astype(np.float32)
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
ann = ANNFactory.create({"backend": "ggml", "ggml": {"quantize": "NOEXIST", "gpu": False}})
|
||||
ann.index(data)
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
ann = ANNFactory.create({"backend": "ggml", "ggml": {"quantize": "Q4_K"}})
|
||||
ann.index(data)
|
||||
|
||||
def testHnsw(self):
|
||||
"""
|
||||
Test Hnswlib backend
|
||||
"""
|
||||
|
||||
self.runTests("hnsw")
|
||||
|
||||
def testHnswCustom(self):
|
||||
"""
|
||||
Test Hnswlib backend with custom settings
|
||||
"""
|
||||
|
||||
# Test with custom settings
|
||||
self.runTests("hnsw", {"hnsw": {"efconstruction": 100, "m": 4, "randomseed": 0, "efsearch": 5}})
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
ann = ANN({})
|
||||
|
||||
self.assertRaises(NotImplementedError, ann.load, None)
|
||||
self.assertRaises(NotImplementedError, ann.index, None)
|
||||
self.assertRaises(NotImplementedError, ann.append, None)
|
||||
self.assertRaises(NotImplementedError, ann.delete, None)
|
||||
self.assertRaises(NotImplementedError, ann.search, None, None)
|
||||
self.assertRaises(NotImplementedError, ann.count)
|
||||
self.assertRaises(NotImplementedError, ann.save, None)
|
||||
|
||||
def testNumPy(self):
|
||||
"""
|
||||
Test NumPy backend
|
||||
"""
|
||||
|
||||
self.runTests("numpy")
|
||||
|
||||
@patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
|
||||
def testNumPyLegacy(self):
|
||||
"""
|
||||
Test NumPy backend with legacy pickled data
|
||||
"""
|
||||
|
||||
serializer = SerializeFactory.create("pickle", allowpickle=True)
|
||||
|
||||
# Create output directory
|
||||
output = os.path.join(tempfile.gettempdir(), "ann.npy")
|
||||
path = os.path.join(output, "embeddings")
|
||||
os.makedirs(output, exist_ok=True)
|
||||
|
||||
# Generate data and save as pickle
|
||||
data = np.random.rand(100, 240).astype(np.float32)
|
||||
serializer.save(data, path)
|
||||
|
||||
ann = ANNFactory.create({"backend": "numpy"})
|
||||
ann.load(path)
|
||||
|
||||
# Validate count
|
||||
self.assertEqual(ann.count(), 100)
|
||||
|
||||
def testNumPySafetensors(self):
|
||||
"""
|
||||
Test NumPy backend with safetensors storage
|
||||
"""
|
||||
|
||||
ann = ANNFactory.create({"backend": "numpy", "numpy": {"safetensors": True}})
|
||||
|
||||
# Generate and index dummy data
|
||||
data = np.random.rand(100, 240).astype(np.float32)
|
||||
ann.index(data)
|
||||
|
||||
# Test save and load
|
||||
index = os.path.join(tempfile.gettempdir(), "numpy.safetensors")
|
||||
ann.save(index)
|
||||
ann.load(index)
|
||||
|
||||
# Generate query vector and test search
|
||||
query = np.random.rand(240).astype(np.float32)
|
||||
self.normalize(query)
|
||||
self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)
|
||||
|
||||
# Validate count
|
||||
self.assertEqual(ann.count(), 100)
|
||||
|
||||
@patch("sqlalchemy.orm.Query.limit")
|
||||
def testPGVector(self, query):
|
||||
"""
|
||||
Test PGVector backend
|
||||
"""
|
||||
|
||||
# Generate test record
|
||||
data = np.random.rand(1, 240).astype(np.float32)
|
||||
|
||||
# Mock database query
|
||||
query.return_value = [(x, -1.0) for x in range(data.shape[0])]
|
||||
|
||||
configs = [
|
||||
("full", {"dimensions": 240}, {}, data),
|
||||
("half", {"dimensions": 240}, {"precision": "half"}, data),
|
||||
("binary", {"quantize": 1, "dimensions": 240 * 8}, {}, data.astype(np.uint8)),
|
||||
]
|
||||
|
||||
# Create ANN
|
||||
for name, config, pgvector, data in configs:
|
||||
path = os.path.join(tempfile.gettempdir(), f"pgvector.{name}.sqlite")
|
||||
ann = ANNFactory.create(
|
||||
{**{"backend": "pgvector", "pgvector": {**{"url": f"sqlite:///{path}", "schema": "txtai"}, **pgvector}}, **config}
|
||||
)
|
||||
|
||||
# Test indexing
|
||||
ann.index(data)
|
||||
ann.append(data)
|
||||
|
||||
# Validate search results
|
||||
self.assertEqual(ann.search(data, 1), [[(0, 1.0)]])
|
||||
|
||||
# Validate save/load/delete
|
||||
ann.save(None)
|
||||
ann.load(None)
|
||||
|
||||
# Validate count
|
||||
self.assertEqual(ann.count(), 2)
|
||||
|
||||
# Test delete
|
||||
ann.delete([0])
|
||||
self.assertEqual(ann.count(), 1)
|
||||
|
||||
# Close ANN
|
||||
ann.close()
|
||||
|
||||
@unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS")
|
||||
def testSQLite(self):
|
||||
"""
|
||||
Test SQLite backend
|
||||
"""
|
||||
|
||||
self.runTests("sqlite")
|
||||
|
||||
@unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS")
|
||||
def testSQLiteCustom(self):
|
||||
"""
|
||||
Test SQLite backend with custom settings
|
||||
"""
|
||||
|
||||
# Test with custom settings
|
||||
self.runTests("sqlite", {"sqlite": {"quantize": 1}})
|
||||
self.runTests("sqlite", {"sqlite": {"quantize": 8}})
|
||||
|
||||
# Test saving to a new path
|
||||
model = self.backend("sqlite")
|
||||
expected = model.count() - 1
|
||||
|
||||
# Test save variations
|
||||
index = os.path.join(tempfile.gettempdir(), "ann.sqlite")
|
||||
new = os.path.join(tempfile.gettempdir(), "ann.sqlite.new")
|
||||
|
||||
# Save new
|
||||
model.save(index)
|
||||
|
||||
# Save to same path
|
||||
model.save(index)
|
||||
|
||||
# Delete id
|
||||
model.delete([0])
|
||||
|
||||
# Save to another path
|
||||
model.load(index)
|
||||
model.save(new)
|
||||
|
||||
self.assertEqual(model.count(), expected)
|
||||
|
||||
def testTorch(self):
|
||||
"""
|
||||
Test Torch backend
|
||||
"""
|
||||
|
||||
self.runTests("torch")
|
||||
|
||||
@unittest.skipIf(platform.system() == "Darwin", "Torch quantization not supported on macOS")
|
||||
def testTorchQuantization(self):
|
||||
"""
|
||||
Test Torch backend with quantization enabled
|
||||
"""
|
||||
|
||||
for qtype in ["fp4", "nf4", "int8"]:
|
||||
ann = ANNFactory.create({"backend": "torch", "torch": {"quantize": {"type": qtype}}})
|
||||
|
||||
# Generate and index dummy data
|
||||
data = np.random.rand(100, 240).astype(np.float32)
|
||||
ann.index(data)
|
||||
|
||||
# Test save and load
|
||||
index = os.path.join(tempfile.gettempdir(), f"{qtype}.safetensors")
|
||||
ann.save(index)
|
||||
ann.load(index)
|
||||
|
||||
# Generate query vector and test search
|
||||
query = np.random.rand(240).astype(np.float32)
|
||||
self.normalize(query)
|
||||
self.assertGreater(ann.search(np.array([query]), 1)[0][0][1], 0)
|
||||
|
||||
# Validate count
|
||||
self.assertEqual(ann.count(), 100)
|
||||
|
||||
# Test delete
|
||||
ann.delete([0])
|
||||
self.assertEqual(ann.count(), 99)
|
||||
|
||||
def testTurboVec(self):
|
||||
"""
|
||||
Test turbovec backend
|
||||
"""
|
||||
|
||||
self.runTests("turbovec")
|
||||
|
||||
def runTests(self, name, params=None, update=True):
|
||||
"""
|
||||
Runs a series of standard backend tests.
|
||||
|
||||
Args:
|
||||
name: backend name
|
||||
params: additional config parameters
|
||||
update: If append/delete options should be tested
|
||||
"""
|
||||
|
||||
self.assertEqual(self.backend(name, params).config["backend"], name)
|
||||
self.assertEqual(self.save(name, params).count(), 10000)
|
||||
|
||||
if update:
|
||||
self.assertEqual(self.append(name, params, 500).count(), 10500)
|
||||
self.assertEqual(self.delete(name, params, [0, 1]).count(), 9998)
|
||||
self.assertEqual(self.delete(name, params, [100000]).count(), 10000)
|
||||
|
||||
self.assertGreater(self.search(name, params), 0)
|
||||
|
||||
def backend(self, name, params=None, length=10000):
|
||||
"""
|
||||
Test a backend.
|
||||
|
||||
Args:
|
||||
name: backend name
|
||||
params: additional config parameters
|
||||
length: number of rows to generate
|
||||
|
||||
Returns:
|
||||
ANN model
|
||||
"""
|
||||
|
||||
# Generate test data
|
||||
data = np.random.rand(length, 240).astype(np.float32)
|
||||
self.normalize(data)
|
||||
|
||||
config = {"backend": name, "dimensions": data.shape[1]}
|
||||
if params:
|
||||
config.update(params)
|
||||
|
||||
model = ANNFactory.create(config)
|
||||
model.index(data)
|
||||
|
||||
return model
|
||||
|
||||
def append(self, name, params=None, length=500):
|
||||
"""
|
||||
Appends new data to index.
|
||||
|
||||
Args:
|
||||
name: backend name
|
||||
params: additional config parameters
|
||||
length: number of rows to generate
|
||||
|
||||
Returns:
|
||||
ANN model
|
||||
"""
|
||||
|
||||
# Initial model
|
||||
model = self.backend(name, params)
|
||||
|
||||
# Generate test data
|
||||
data = np.random.rand(length, 240).astype(np.float32)
|
||||
self.normalize(data)
|
||||
|
||||
model.append(data)
|
||||
|
||||
return model
|
||||
|
||||
def delete(self, name, params=None, ids=None):
|
||||
"""
|
||||
Deletes data from index.
|
||||
|
||||
Args:
|
||||
name: backend name
|
||||
params: additional config parameters
|
||||
ids: ids to delete
|
||||
|
||||
Returns:
|
||||
ANN model
|
||||
"""
|
||||
|
||||
# Initial model
|
||||
model = self.backend(name, params)
|
||||
model.delete(ids)
|
||||
|
||||
return model
|
||||
|
||||
def save(self, name, params=None):
|
||||
"""
|
||||
Test save/load.
|
||||
|
||||
Args:
|
||||
name: backend name
|
||||
params: additional config parameters
|
||||
|
||||
Returns:
|
||||
ANN model
|
||||
"""
|
||||
|
||||
model = self.backend(name, params)
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "ann")
|
||||
|
||||
# Save and close index
|
||||
model.save(index)
|
||||
model.close()
|
||||
|
||||
# Reload index
|
||||
model.load(index)
|
||||
|
||||
return model
|
||||
|
||||
def search(self, name, params=None):
|
||||
"""
|
||||
Test ANN search.
|
||||
|
||||
Args:
|
||||
name: backend name
|
||||
params: additional config parameters
|
||||
|
||||
Returns:
|
||||
search results
|
||||
"""
|
||||
|
||||
# Generate ANN index
|
||||
model = self.backend(name, params)
|
||||
|
||||
# Generate query vector
|
||||
query = np.random.rand(240).astype(np.float32)
|
||||
self.normalize(query)
|
||||
|
||||
# Ensure top result has similarity > 0
|
||||
return model.search(np.array([query]), 1)[0][0][1]
|
||||
|
||||
def normalize(self, embeddings):
|
||||
"""
|
||||
Normalizes embeddings using L2 normalization. Operation applied directly on array.
|
||||
|
||||
Args:
|
||||
embeddings: input embeddings matrix
|
||||
"""
|
||||
|
||||
# Calculation is different for matrices vs vectors
|
||||
if len(embeddings.shape) > 1:
|
||||
embeddings /= np.linalg.norm(embeddings, axis=1)[:, np.newaxis]
|
||||
else:
|
||||
embeddings /= np.linalg.norm(embeddings)
|
||||
@@ -0,0 +1,165 @@
|
||||
"""
|
||||
Sparse ANN module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from scipy.sparse import random
|
||||
from sklearn.preprocessing import normalize
|
||||
|
||||
from txtai.ann import SparseANNFactory
|
||||
|
||||
|
||||
class TestSparse(unittest.TestCase):
|
||||
"""
|
||||
Sparse ANN tests.
|
||||
"""
|
||||
|
||||
def testCustomBackend(self):
|
||||
"""
|
||||
Test resolving a custom backend
|
||||
"""
|
||||
|
||||
self.assertIsNotNone(SparseANNFactory.create({"backend": "txtai.ann.IVFSparse"}))
|
||||
|
||||
def testCustomBackendNotFound(self):
|
||||
"""
|
||||
Test resolving an unresolvable backend
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
SparseANNFactory.create({"backend": "notfound.ann"})
|
||||
|
||||
def testIVFSparse(self):
|
||||
"""
|
||||
Test IVFSparse backend
|
||||
"""
|
||||
|
||||
# Generate test record
|
||||
insert = self.generate(500, 30522)
|
||||
append = self.generate(500, 30522)
|
||||
|
||||
# Count of records
|
||||
count = insert.shape[0] + append.shape[0]
|
||||
|
||||
# Create ANN
|
||||
path = os.path.join(tempfile.gettempdir(), "ivfsparse")
|
||||
ann = SparseANNFactory.create({"backend": "ivfsparse", "ivfsparse": {"nlist": 2, "nprobe": 2, "sample": 1.0}})
|
||||
|
||||
# Test indexing
|
||||
ann.index(insert)
|
||||
ann.append(append)
|
||||
|
||||
# Validate search results
|
||||
results = [x[0] for x in ann.search(insert[5], 10)[0]]
|
||||
self.assertIn(5, results)
|
||||
|
||||
# Validate save/load/delete
|
||||
ann.save(path)
|
||||
ann.load(path)
|
||||
|
||||
# Validate count
|
||||
self.assertEqual(ann.count(), count)
|
||||
|
||||
# Test delete
|
||||
ann.delete([0])
|
||||
self.assertEqual(ann.count(), count - 1)
|
||||
|
||||
# Re-validate search results
|
||||
results = [x[0] for x in ann.search(append[0], 10)[0]]
|
||||
self.assertIn(insert.shape[0], results)
|
||||
|
||||
# Close ANN
|
||||
ann.close()
|
||||
|
||||
# Test cluster pruning
|
||||
ann = SparseANNFactory.create({"backend": "ivfsparse", "ivfsparse": {"nlist": 15, "nprobe": 1, "sample": 1.0}})
|
||||
ann.index(insert)
|
||||
self.assertLess(len(ann.blocks), 15)
|
||||
ann.close()
|
||||
|
||||
def testIVFSparseTopnOverLimit(self):
|
||||
"""
|
||||
Test IVFSparse topn when limit exceeds the number of indexed documents
|
||||
"""
|
||||
|
||||
# Generate a small dataset (5 documents)
|
||||
data = self.generate(5, 30522)
|
||||
|
||||
ann = SparseANNFactory.create({"backend": "ivfsparse"})
|
||||
ann.index(data)
|
||||
|
||||
# Search with limit (10) greater than document count (5)
|
||||
results = ann.search(data[0], 10)
|
||||
self.assertGreater(len(results[0]), 0)
|
||||
|
||||
# Batch search with multiple queries exceeding document count
|
||||
results = ann.search(data, 10)
|
||||
self.assertEqual(len(results), data.shape[0])
|
||||
for result in results:
|
||||
self.assertGreater(len(result), 0)
|
||||
|
||||
ann.close()
|
||||
|
||||
@patch("sqlalchemy.orm.Query.limit")
|
||||
def testPGSparse(self, query):
|
||||
"""
|
||||
Test Sparse Postgres backend
|
||||
"""
|
||||
|
||||
# Generate test record
|
||||
data = self.generate(1, 30522)
|
||||
|
||||
# Mock database query
|
||||
query.return_value = [(x, -1.0) for x in range(data.shape[0])]
|
||||
|
||||
# Create ANN
|
||||
path = os.path.join(tempfile.gettempdir(), "pgsparse.sqlite")
|
||||
ann = SparseANNFactory.create({"backend": "pgsparse", "dimensions": 30522, "pgsparse": {"url": f"sqlite:///{path}", "schema": "txtai"}})
|
||||
|
||||
# Test indexing
|
||||
ann.index(data)
|
||||
ann.append(data)
|
||||
|
||||
# Validate search results
|
||||
self.assertEqual(ann.search(data, 1), [[(0, 1.0)]])
|
||||
|
||||
# Validate save/load/delete
|
||||
ann.save(None)
|
||||
ann.load(None)
|
||||
|
||||
# Validate count
|
||||
self.assertEqual(ann.count(), 2)
|
||||
|
||||
# Test delete
|
||||
ann.delete([0])
|
||||
self.assertEqual(ann.count(), 1)
|
||||
|
||||
# Test > 1000 dimensions
|
||||
data = random(1, 30522, format="csr", density=0.1)
|
||||
ann.index(data)
|
||||
self.assertEqual(ann.count(), 1)
|
||||
|
||||
# Close ANN
|
||||
ann.close()
|
||||
|
||||
def generate(self, m, n):
|
||||
"""
|
||||
Generates random normalized sparse data.
|
||||
|
||||
Args:
|
||||
m, n: shape of the matrix
|
||||
|
||||
Returns:
|
||||
csr matrix
|
||||
"""
|
||||
|
||||
# Generate random csr matrix
|
||||
data = random(m, n, format="csr")
|
||||
|
||||
# Normalize and return
|
||||
return normalize(data)
|
||||
@@ -0,0 +1,91 @@
|
||||
"""
|
||||
Agent API module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import API, application
|
||||
|
||||
# Configuration for agents
|
||||
AGENTS = """
|
||||
agent:
|
||||
test:
|
||||
max_iterations: 1
|
||||
tools:
|
||||
- name: testtool
|
||||
description: Test tool
|
||||
target: testapi.testapiagent.TestTool
|
||||
|
||||
llm:
|
||||
path: hf-internal-testing/tiny-random-LlamaForCausalLM
|
||||
"""
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestAgent(unittest.TestCase):
|
||||
"""
|
||||
API tests for agents.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
|
||||
def start():
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testapi.yml")
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write(AGENTS)
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
# Patch LLM to generate answer
|
||||
agent = application.get().agents["test"]
|
||||
agent.process.model.llm = lambda *args, **kwargs: 'Action:\n{"name": "final_answer", "arguments": "Hi"}'
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestAgent.start()
|
||||
|
||||
def testAgent(self):
|
||||
"""
|
||||
Test agent via API
|
||||
"""
|
||||
|
||||
results = self.client.post("agent", json={"name": "test", "text": "Hello"}).json()
|
||||
self.assertEqual(results, "Hi")
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test empty API configuration
|
||||
"""
|
||||
|
||||
api = API({})
|
||||
|
||||
self.assertIsNone(api.agent("junk", "test"))
|
||||
|
||||
|
||||
class TestTool:
|
||||
"""
|
||||
Class to test agent tools
|
||||
"""
|
||||
|
||||
def __call__(self):
|
||||
pass
|
||||
@@ -0,0 +1,502 @@
|
||||
"""
|
||||
Embeddings API module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
import urllib.parse
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import API, application
|
||||
|
||||
# Configuration for a read/write embeddings index
|
||||
INDEX = """
|
||||
# Index file path
|
||||
path: %s
|
||||
|
||||
# Allow indexing of documents
|
||||
writable: True
|
||||
|
||||
# Allow reindexing
|
||||
reindex: True
|
||||
|
||||
# Questions settings
|
||||
questions:
|
||||
path: distilbert-base-cased-distilled-squad
|
||||
|
||||
# Embeddings settings
|
||||
embeddings:
|
||||
path: sentence-transformers/nli-mpnet-base-v2
|
||||
|
||||
# Extractor settings
|
||||
extractor:
|
||||
path: questions
|
||||
"""
|
||||
|
||||
# Configuration for a read-only embeddings index
|
||||
READONLY = """
|
||||
# Index file path
|
||||
path: %s
|
||||
|
||||
# Allow indexing of documents
|
||||
writable: False
|
||||
|
||||
# Embeddings settings
|
||||
embeddings:
|
||||
"""
|
||||
|
||||
# Configuration for an index with custom functions
|
||||
FUNCTIONS = """
|
||||
# Ignore existing index
|
||||
pathignore: %s
|
||||
|
||||
# Allow indexing of documents
|
||||
writable: True
|
||||
|
||||
# Embeddings settings
|
||||
embeddings:
|
||||
path: sentence-transformers/nli-mpnet-base-v2
|
||||
content: True
|
||||
functions:
|
||||
- testapi.testapiembeddings.Elements
|
||||
- name: length
|
||||
argcount: 1
|
||||
function: testapi.testapiembeddings.length
|
||||
- name: ann
|
||||
function: ann
|
||||
transform: testapi.testapiembeddings.transform
|
||||
"""
|
||||
|
||||
# Configuration for RAG
|
||||
RAG = """
|
||||
# Ignore existing index
|
||||
pathignore: %s
|
||||
|
||||
# Allow indexing of documents
|
||||
writable: True
|
||||
|
||||
# Embeddings settings
|
||||
embeddings:
|
||||
path: sentence-transformers/nli-mpnet-base-v2
|
||||
content: True
|
||||
|
||||
# LLM
|
||||
llm:
|
||||
path: hf-internal-testing/tiny-random-gpt2
|
||||
task: language-generation
|
||||
|
||||
# RAG settings
|
||||
rag:
|
||||
path: llm
|
||||
output: flatten
|
||||
"""
|
||||
|
||||
# Configuration for reindexing disabled
|
||||
REINDEXDISABLED = """
|
||||
# Index file path
|
||||
path: %s
|
||||
|
||||
# Allow indexing of documents
|
||||
writable: True
|
||||
|
||||
# Embeddings settings
|
||||
embeddings:
|
||||
path: sentence-transformers/nli-mpnet-base-v2
|
||||
content: True
|
||||
"""
|
||||
|
||||
# Configuration for reranker
|
||||
RERANK = """
|
||||
# Index file path
|
||||
path: %s
|
||||
|
||||
# Allow indexing of documents
|
||||
writable: True
|
||||
|
||||
# Embeddings settings
|
||||
embeddings:
|
||||
path: sentence-transformers/nli-mpnet-base-v2
|
||||
content: True
|
||||
|
||||
# Similarity and Reranking settings
|
||||
similarity:
|
||||
path: neuml/colbert-bert-tiny
|
||||
lateencode: True
|
||||
|
||||
reranker:
|
||||
"""
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestEmbeddings(unittest.TestCase):
|
||||
"""
|
||||
API tests for embeddings indices.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
|
||||
def start(yaml, path="testapi"):
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
|
||||
Args:
|
||||
yaml: input configuration
|
||||
path: output path
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testapi.yml")
|
||||
index = os.path.join(tempfile.gettempdir(), path)
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write(yaml % index)
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestEmbeddings.start(INDEX, "testapi")
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Index data
|
||||
cls.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(cls.data)])
|
||||
cls.client.get("index")
|
||||
|
||||
def testCount(self):
|
||||
"""
|
||||
Test count via API
|
||||
"""
|
||||
|
||||
self.assertEqual(self.client.get("count").json(), 6)
|
||||
|
||||
def testDelete(self):
|
||||
"""
|
||||
Test delete via API
|
||||
"""
|
||||
|
||||
# Delete best match
|
||||
ids = self.client.post("delete", json=[4]).json()
|
||||
self.assertEqual(ids, [4])
|
||||
|
||||
# Search for best match
|
||||
query = urllib.parse.quote("feel good story")
|
||||
uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
|
||||
|
||||
self.assertEqual(self.client.get("count").json(), 5)
|
||||
self.assertEqual(uid, 5)
|
||||
|
||||
# Reset data
|
||||
self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
|
||||
self.client.get("index")
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test empty API configuration
|
||||
"""
|
||||
|
||||
api = API({"writable": True})
|
||||
|
||||
self.assertIsNone(api.search("test", None))
|
||||
self.assertIsNone(api.batchsearch(["test"], None))
|
||||
self.assertIsNone(api.delete(["test"]))
|
||||
self.assertIsNone(api.count())
|
||||
self.assertIsNone(api.similarity("test", ["test"]))
|
||||
self.assertIsNone(api.batchsimilarity(["test"], ["test"]))
|
||||
self.assertIsNone(api.explain("test"))
|
||||
self.assertIsNone(api.batchexplain(["test"]))
|
||||
self.assertIsNone(api.transform("test"))
|
||||
self.assertIsNone(api.batchtransform(["test"]))
|
||||
self.assertIsNone(api.extract(["test"], ["test"]))
|
||||
|
||||
def testExtractor(self):
|
||||
"""
|
||||
Test qa extraction via API
|
||||
"""
|
||||
|
||||
data = [
|
||||
"Giants hit 3 HRs to down Dodgers",
|
||||
"Giants 5 Dodgers 4 final",
|
||||
"Dodgers drop Game 2 against the Giants, 5-4",
|
||||
"Blue Jays beat Red Sox final score 2-1",
|
||||
"Red Sox lost to the Blue Jays, 2-1",
|
||||
"Blue Jays at Red Sox is over. Score: 2-1",
|
||||
"Phillies win over the Braves, 5-0",
|
||||
"Phillies 5 Braves 0 final",
|
||||
"Final: Braves lose to the Phillies in the series opener, 5-0",
|
||||
"Lightning goaltender pulled, lose to Flyers 4-1",
|
||||
"Flyers 4 Lightning 1 final",
|
||||
"Flyers win 4-1",
|
||||
]
|
||||
|
||||
questions = ["What team won the game?", "What was score?"]
|
||||
|
||||
# pylint: disable=C3001
|
||||
execute = lambda query: self.client.post(
|
||||
"extract",
|
||||
json={"queue": [{"name": question, "query": query, "question": question, "snippet": False} for question in questions], "texts": data},
|
||||
).json()
|
||||
|
||||
answers = execute("Red Sox - Blue Jays")
|
||||
self.assertEqual("Blue Jays", answers[0]["answer"])
|
||||
self.assertEqual("2-1", answers[1]["answer"])
|
||||
|
||||
# Ad-hoc questions
|
||||
question = "What hockey team won?"
|
||||
|
||||
answers = self.client.post(
|
||||
"extract", json={"queue": [{"name": question, "query": question, "question": question, "snippet": False}], "texts": data}
|
||||
).json()
|
||||
self.assertEqual("Flyers", answers[0]["answer"])
|
||||
|
||||
def testReindex(self):
|
||||
"""
|
||||
Test reindex via API
|
||||
"""
|
||||
|
||||
# Reindex data
|
||||
self.client.post("reindex", json={"config": {"path": "sentence-transformers/nli-mpnet-base-v2"}})
|
||||
|
||||
# Search for best match
|
||||
query = urllib.parse.quote("feel good story")
|
||||
uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
|
||||
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Reset data
|
||||
self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
|
||||
self.client.get("index")
|
||||
|
||||
def testSearch(self):
|
||||
"""
|
||||
Test search via API
|
||||
"""
|
||||
|
||||
query = urllib.parse.quote("feel good story")
|
||||
uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testSearchBatch(self):
|
||||
"""
|
||||
Test batch search via API
|
||||
"""
|
||||
|
||||
results = self.client.post("batchsearch", json={"queries": ["feel good story", "climate change"], "limit": 1}).json()
|
||||
|
||||
uids = [result[0]["id"] for result in results]
|
||||
self.assertEqual(uids, [4, 1])
|
||||
|
||||
def testSimilarity(self):
|
||||
"""
|
||||
Test similarity via API
|
||||
"""
|
||||
|
||||
uid = self.client.post("similarity", json={"query": "feel good story", "texts": self.data}).json()[0]["id"]
|
||||
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testSimilarityBatch(self):
|
||||
"""
|
||||
Test batch similarity via API
|
||||
"""
|
||||
|
||||
results = self.client.post("batchsimilarity", json={"queries": ["feel good story", "climate change"], "texts": self.data}).json()
|
||||
|
||||
uids = [result[0]["id"] for result in results]
|
||||
self.assertEqual(uids, [4, 1])
|
||||
|
||||
def testTransform(self):
|
||||
"""
|
||||
Test embeddings transform via API
|
||||
"""
|
||||
|
||||
self.assertEqual(len(self.client.get("transform?text=testembed").json()), 768)
|
||||
|
||||
def testTransformBatch(self):
|
||||
"""
|
||||
Test batch embeddings transform via API
|
||||
"""
|
||||
|
||||
embeddings = self.client.post("batchtransform", json=self.data).json()
|
||||
|
||||
self.assertEqual(len(embeddings), len(self.data))
|
||||
self.assertEqual(len(embeddings[0]), 768)
|
||||
|
||||
def testUpsert(self):
|
||||
"""
|
||||
Test upsert via API
|
||||
"""
|
||||
|
||||
# Update data
|
||||
self.client.post("add", json=[{"id": 0, "text": "Feel good story: baby panda born"}])
|
||||
self.client.get("upsert")
|
||||
|
||||
# Search for best match
|
||||
query = urllib.parse.quote("feel good story")
|
||||
uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
|
||||
|
||||
self.assertEqual(uid, 0)
|
||||
|
||||
# Reset data
|
||||
self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
|
||||
self.client.get("index")
|
||||
|
||||
def testViewOnly(self):
|
||||
"""
|
||||
Test read-only API instance
|
||||
"""
|
||||
|
||||
# Re-create application with a read-only index
|
||||
self.client = TestEmbeddings.start(READONLY)
|
||||
|
||||
# Test search
|
||||
query = urllib.parse.quote("feel good story")
|
||||
uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Test similarity
|
||||
uid = self.client.post("similarity", json={"query": "feel good story", "texts": self.data}).json()[0]["id"]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Test errors raised for write operations
|
||||
self.assertEqual(self.client.post("add", json=[{"id": 0, "text": "test"}]).status_code, 403)
|
||||
self.assertEqual(self.client.get("index").status_code, 403)
|
||||
self.assertEqual(self.client.get("upsert").status_code, 403)
|
||||
self.assertEqual(self.client.post("delete", json=[0]).status_code, 403)
|
||||
self.assertEqual(self.client.post("reindex", json={"config": {"path": "sentence-transformers/nli-mpnet-base-v2"}}).status_code, 403)
|
||||
|
||||
def testXFunctions(self):
|
||||
"""
|
||||
Test API instance with custom functions
|
||||
"""
|
||||
|
||||
# Re-create application with custom functions
|
||||
self.client = TestEmbeddings.start(FUNCTIONS)
|
||||
|
||||
# Index data
|
||||
self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
|
||||
self.client.get("index")
|
||||
|
||||
query = urllib.parse.quote("select elements('text') length from txtai limit 1")
|
||||
self.assertEqual(self.client.get(f"search?query={query}").json()[0]["length"], 4)
|
||||
|
||||
query = urllib.parse.quote("select length('text') length from txtai limit 1")
|
||||
self.assertEqual(self.client.get(f"search?query={query}").json()[0]["length"], 4)
|
||||
|
||||
def testXPlain(self):
|
||||
"""
|
||||
Test API instance with explain methods
|
||||
"""
|
||||
|
||||
results = self.client.post("explain", json={"query": "feel good story", "limit": 1}).json()
|
||||
|
||||
self.assertEqual(results[0]["text"], self.data[4])
|
||||
self.assertIsNotNone(results[0].get("tokens"))
|
||||
|
||||
def testXPlainBatch(self):
|
||||
"""
|
||||
Test batch query explain via API
|
||||
"""
|
||||
|
||||
results = self.client.post("batchexplain", json={"queries": ["feel good story", "climate change"], "limit": 1}).json()
|
||||
|
||||
text = [result[0]["text"] for result in results]
|
||||
self.assertEqual(text, [self.data[4], self.data[1]])
|
||||
self.assertIsNotNone(results[0][0].get("tokens"))
|
||||
|
||||
def testXRAG(self):
|
||||
"""
|
||||
Test RAG via API
|
||||
"""
|
||||
|
||||
# Re-create application with a RAG pipeline
|
||||
self.client = TestEmbeddings.start(RAG)
|
||||
|
||||
# Index data
|
||||
self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
|
||||
self.client.get("index")
|
||||
|
||||
response = self.client.get("rag?query=bear").json()
|
||||
self.assertIsInstance(response, str)
|
||||
|
||||
response = self.client.post("batchrag", json={"queries": ["bear", "bear"]}).json()
|
||||
self.assertEqual(len(response), 2)
|
||||
|
||||
def testXReindexDisabled(self):
|
||||
"""
|
||||
Test reindexing is disabled
|
||||
"""
|
||||
|
||||
# Re-create application with reindexing disabled
|
||||
self.client = TestEmbeddings.start(REINDEXDISABLED, "testapi-reindex")
|
||||
|
||||
# Index data
|
||||
self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
|
||||
self.client.get("index")
|
||||
|
||||
# Assert error raised
|
||||
self.assertEqual(self.client.post("reindex", json={"config": {"path": "sentence-transformers/nli-mpnet-base-v2"}}).status_code, 403)
|
||||
|
||||
def testXRerank(self):
|
||||
"""
|
||||
Test rerank via API
|
||||
"""
|
||||
|
||||
# Re-create application with a reranker pipeline
|
||||
self.client = TestEmbeddings.start(RERANK, "testapi-rerank")
|
||||
|
||||
# Index data
|
||||
self.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(self.data)])
|
||||
self.client.get("index")
|
||||
|
||||
uid = self.client.get("rerank?query=bear").json()[0]["id"]
|
||||
self.assertEqual(uid, "3")
|
||||
|
||||
results = self.client.post("batchrerank", json={"queries": ["bear", "bear"]}).json()
|
||||
|
||||
uids = [result[0]["id"] for result in results]
|
||||
self.assertEqual(uids, ["3", "3"])
|
||||
|
||||
|
||||
class Elements:
|
||||
"""
|
||||
Custom SQL function as callable object.
|
||||
"""
|
||||
|
||||
def __call__(self, text):
|
||||
return length(text)
|
||||
|
||||
|
||||
def transform(document):
|
||||
"""
|
||||
Custom transform function.
|
||||
"""
|
||||
|
||||
return document
|
||||
|
||||
|
||||
def length(text):
|
||||
"""
|
||||
Custom SQL function.
|
||||
"""
|
||||
|
||||
return len(text)
|
||||
@@ -0,0 +1,391 @@
|
||||
"""
|
||||
Pipeline API module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
import urllib
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import API, application
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
# Configuration for pipelines
|
||||
PIPELINES = """
|
||||
# Image captions
|
||||
caption:
|
||||
|
||||
# Entity extraction
|
||||
entity:
|
||||
path: dslim/bert-base-NER
|
||||
|
||||
# Extractor settings
|
||||
extractor:
|
||||
similarity: similarity
|
||||
path: llm
|
||||
|
||||
# Label settings
|
||||
labels:
|
||||
path: prajjwal1/bert-medium-mnli
|
||||
|
||||
# LLM settings
|
||||
llm:
|
||||
path: hf-internal-testing/tiny-random-gpt2
|
||||
task: language-generation
|
||||
|
||||
# Image objects
|
||||
objects:
|
||||
|
||||
# Text segmentation
|
||||
segmentation:
|
||||
sentences: true
|
||||
|
||||
# Enable pipeline similarity backed by zero shot classifier
|
||||
similarity:
|
||||
|
||||
# Summarization
|
||||
summary:
|
||||
path: t5-small
|
||||
|
||||
# Tabular
|
||||
tabular:
|
||||
|
||||
# Text extraction
|
||||
textractor:
|
||||
safeopen: /tmp/txtai
|
||||
|
||||
# Text to speech
|
||||
texttospeech:
|
||||
|
||||
# Transcription
|
||||
transcription:
|
||||
|
||||
# Translation:
|
||||
translation:
|
||||
|
||||
# Enable file uploads
|
||||
upload:
|
||||
"""
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestPipeline(unittest.TestCase):
|
||||
"""
|
||||
API tests for pipelines.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
|
||||
def start():
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testapi.yml")
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write(PIPELINES)
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestPipeline.start()
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
cls.text = (
|
||||
"Search is the base of many applications. Once data starts to pile up, users want to be able to find it. It's the foundation "
|
||||
"of the internet and an ever-growing challenge that is never solved or done. The field of Natural Language Processing (NLP) is "
|
||||
"rapidly evolving with a number of new developments. Large-scale general language models are an exciting new capability "
|
||||
"allowing us to add amazing functionality quickly with limited compute and people. Innovation continues with new models "
|
||||
"and advancements coming in at what seems a weekly basis. This article introduces txtai, an AI-powered search engine "
|
||||
"that enables Natural Language Understanding (NLU) based search in any application."
|
||||
)
|
||||
|
||||
# Create invalid test file
|
||||
directory = os.path.join(tempfile.gettempdir(), "txtai-1")
|
||||
os.makedirs(directory, exist_ok=True)
|
||||
with open(os.path.join(directory, "test"), "w", encoding="utf-8") as output:
|
||||
output.write("123")
|
||||
|
||||
def testCaption(self):
|
||||
"""
|
||||
Test caption via API
|
||||
"""
|
||||
|
||||
caption = self.client.get(f"caption?file={Utils.PATH}/books.jpg").json()
|
||||
|
||||
self.assertEqual(caption, "a book shelf filled with books and a stack of books")
|
||||
|
||||
def testCaptionBatch(self):
|
||||
"""
|
||||
Test batch caption via API
|
||||
"""
|
||||
|
||||
path = Utils.PATH + "/books.jpg"
|
||||
|
||||
captions = self.client.post("batchcaption", json=[path, path]).json()
|
||||
self.assertEqual(captions, ["a book shelf filled with books and a stack of books"] * 2)
|
||||
|
||||
def testEntity(self):
|
||||
"""
|
||||
Test entity extraction via API
|
||||
"""
|
||||
|
||||
entities = self.client.get(f"entity?text={self.data[1]}").json()
|
||||
self.assertEqual([e[0] for e in entities], ["Canada", "Manhattan"])
|
||||
|
||||
def testEntityBatch(self):
|
||||
"""
|
||||
Test batch entity via API
|
||||
"""
|
||||
|
||||
entities = self.client.post("batchentity", json=[self.data[1]]).json()
|
||||
self.assertEqual([e[0] for e in entities[0]], ["Canada", "Manhattan"])
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test empty API configuration
|
||||
"""
|
||||
|
||||
api = API({})
|
||||
|
||||
self.assertIsNone(api.label("test", ["test"]))
|
||||
self.assertIsNone(api.pipeline("junk", "test"))
|
||||
|
||||
def testLabel(self):
|
||||
"""
|
||||
Test label via API
|
||||
"""
|
||||
|
||||
labels = self.client.post("label", json={"text": "this is the best sentence ever", "labels": ["positive", "negative"]}).json()
|
||||
|
||||
self.assertEqual(labels[0]["id"], 0)
|
||||
|
||||
def testLabelBatch(self):
|
||||
"""
|
||||
Test batch label via API
|
||||
"""
|
||||
|
||||
labels = self.client.post(
|
||||
"batchlabel", json={"texts": ["this is the best sentence ever", "This is terrible"], "labels": ["positive", "negative"]}
|
||||
).json()
|
||||
|
||||
results = [l[0]["id"] for l in labels]
|
||||
self.assertEqual(results, [0, 1])
|
||||
|
||||
def testLLM(self):
|
||||
"""
|
||||
Test LLM inference via API
|
||||
"""
|
||||
|
||||
response = self.client.get("llm?text=test").json()
|
||||
self.assertIsInstance(response, str)
|
||||
|
||||
def testLLMBatch(self):
|
||||
"""
|
||||
Test batch LLM inference via API
|
||||
"""
|
||||
|
||||
response = self.client.post("batchllm", json={"texts": ["test", "test"]}).json()
|
||||
self.assertEqual(len(response), 2)
|
||||
|
||||
def testObjects(self):
|
||||
"""
|
||||
Test objects via API
|
||||
"""
|
||||
|
||||
objects = self.client.get(f"objects?file={Utils.PATH}/books.jpg").json()
|
||||
|
||||
self.assertEqual(objects[0][0], "book")
|
||||
|
||||
def testObjectsBatch(self):
|
||||
"""
|
||||
Test batch objects via API
|
||||
"""
|
||||
|
||||
path = Utils.PATH + "/books.jpg"
|
||||
|
||||
objects = self.client.post("batchobjects", json=[path, path]).json()
|
||||
self.assertEqual([o[0][0] for o in objects], ["book"] * 2)
|
||||
|
||||
def testSegment(self):
|
||||
"""
|
||||
Test segmentation via API
|
||||
"""
|
||||
|
||||
text = self.client.get("segment?text=This is a test. And another test.").json()
|
||||
|
||||
# Check array length is 2
|
||||
self.assertEqual(len(text), 2)
|
||||
|
||||
def testSegmentBatch(self):
|
||||
"""
|
||||
Test batch segmentation via API
|
||||
"""
|
||||
|
||||
text = "This is a test. And another test."
|
||||
texts = self.client.post("batchsegment", json=[text, text]).json()
|
||||
|
||||
# Check array length is 2 and first element length is 2
|
||||
self.assertEqual(len(texts), 2)
|
||||
self.assertEqual(len(texts[0]), 2)
|
||||
|
||||
def testSimilarity(self):
|
||||
"""
|
||||
Test similarity via API
|
||||
"""
|
||||
|
||||
uid = self.client.post("similarity", json={"query": "feel good story", "texts": self.data}).json()[0]["id"]
|
||||
|
||||
self.assertEqual(self.data[uid], self.data[4])
|
||||
|
||||
def testSimilarityBatch(self):
|
||||
"""
|
||||
Test batch similarity via API
|
||||
"""
|
||||
|
||||
results = self.client.post("batchsimilarity", json={"queries": ["feel good story", "climate change"], "texts": self.data}).json()
|
||||
|
||||
uids = [result[0]["id"] for result in results]
|
||||
self.assertEqual(uids, [4, 1])
|
||||
|
||||
def testSummary(self):
|
||||
"""
|
||||
Test summary via API
|
||||
"""
|
||||
|
||||
summary = self.client.get(f"summary?text={urllib.parse.quote(self.text)}&minlength=15&maxlength=15").json()
|
||||
self.assertEqual(summary, "the field of natural language processing (NLP) is rapidly evolving")
|
||||
|
||||
def testSummaryBatch(self):
|
||||
"""
|
||||
Test batch summary via API
|
||||
"""
|
||||
|
||||
summaries = self.client.post("batchsummary", json={"texts": [self.text, self.text], "minlength": 15, "maxlength": 15}).json()
|
||||
self.assertEqual(summaries, ["the field of natural language processing (NLP) is rapidly evolving"] * 2)
|
||||
|
||||
def testTabular(self):
|
||||
"""
|
||||
Test tabular via API
|
||||
"""
|
||||
|
||||
results = self.client.get(f"tabular?file={Utils.PATH}/tabular.csv").json()
|
||||
|
||||
# Check length of results is as expected
|
||||
self.assertEqual(len(results), 6)
|
||||
|
||||
def testTabularBatch(self):
|
||||
"""
|
||||
Test batch tabular via API
|
||||
"""
|
||||
|
||||
path = Utils.PATH + "/tabular.csv"
|
||||
|
||||
results = self.client.post("batchtabular", json=[path, path]).json()
|
||||
self.assertEqual((len(results[0]), len(results[1])), (6, 6))
|
||||
|
||||
def testTextractor(self):
|
||||
"""
|
||||
Test textractor via API
|
||||
"""
|
||||
|
||||
text = self.client.get(f"textract?file={Utils.PATH}/article.pdf").json()
|
||||
|
||||
# Check length of text is as expected
|
||||
self.assertEqual(len(text), 2471)
|
||||
|
||||
# Check invalid URLs
|
||||
for url in ["http://192.168.1.1/path", "http://127.0.0.1/path", "http://invalid", "/etc/config", "/tmp/txtai-1/test"]:
|
||||
with self.assertRaises(IOError):
|
||||
self.client.get(f"textract?file={url}").json()
|
||||
|
||||
def testTextractorBatch(self):
|
||||
"""
|
||||
Test batch textractor via API
|
||||
"""
|
||||
|
||||
path = Utils.PATH + "/article.pdf"
|
||||
|
||||
texts = self.client.post("batchtextract", json=[path, path]).json()
|
||||
self.assertEqual((len(texts[0]), len(texts[1])), (2471, 2471))
|
||||
|
||||
def testTextToSpeech(self):
|
||||
"""
|
||||
Test text to speech
|
||||
"""
|
||||
|
||||
# Generate audio and check for WAV signature
|
||||
audio = self.client.get("texttospeech?text=hello&encoding=wav").content
|
||||
self.assertTrue(audio[0:4] == b"RIFF")
|
||||
|
||||
def testTranscribe(self):
|
||||
"""
|
||||
Test transcribe via API
|
||||
"""
|
||||
|
||||
text = self.client.get(f"transcribe?file={Utils.PATH}/Make_huge_profits.wav").json()
|
||||
|
||||
# Check length of text is as expected
|
||||
self.assertEqual(text, "Make huge profits without working make up to one hundred thousand dollars a day")
|
||||
|
||||
def testTranscribeBatch(self):
|
||||
"""
|
||||
Test batch transcribe via API
|
||||
"""
|
||||
|
||||
path = Utils.PATH + "/Make_huge_profits.wav"
|
||||
|
||||
texts = self.client.post("batchtranscribe", json=[path, path]).json()
|
||||
self.assertEqual(texts, ["Make huge profits without working make up to one hundred thousand dollars a day"] * 2)
|
||||
|
||||
def testTranslate(self):
|
||||
"""
|
||||
Test translate via API
|
||||
"""
|
||||
|
||||
translation = self.client.get(f"translate?text={urllib.parse.quote('This is a test translation into Spanish')}&target=es").json()
|
||||
self.assertEqual(translation, "Esta es una traducción de prueba al español")
|
||||
|
||||
def testTranslateBatch(self):
|
||||
"""
|
||||
Test batch translate via API
|
||||
"""
|
||||
|
||||
text = "This is a test translation into Spanish"
|
||||
translations = self.client.post("batchtranslate", json={"texts": [text, text], "target": "es"}).json()
|
||||
self.assertEqual(translations, ["Esta es una traducción de prueba al español"] * 2)
|
||||
|
||||
def testUpload(self):
|
||||
"""
|
||||
Test file upload
|
||||
"""
|
||||
|
||||
path = Utils.PATH + "/article.pdf"
|
||||
with open(path, "rb") as f:
|
||||
path = self.client.post("upload", files={"files": f}).json()[0]
|
||||
self.assertTrue(os.path.exists(path))
|
||||
@@ -0,0 +1,259 @@
|
||||
"""
|
||||
Workflow API module tests
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from multiprocessing.pool import ThreadPool
|
||||
from threading import Thread
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import API, application
|
||||
|
||||
# Configuration for workflows
|
||||
WORKFLOWS = """
|
||||
# Embeddings index
|
||||
writable: true
|
||||
embeddings:
|
||||
path: sentence-transformers/nli-mpnet-base-v2
|
||||
|
||||
# Labels
|
||||
labels:
|
||||
path: prajjwal1/bert-medium-mnli
|
||||
|
||||
nop:
|
||||
|
||||
# Text segmentation
|
||||
segmentation:
|
||||
sentences: true
|
||||
|
||||
# Workflow definitions
|
||||
workflow:
|
||||
labels:
|
||||
tasks:
|
||||
- action: labels
|
||||
args: [[positive, negative]]
|
||||
multiaction:
|
||||
tasks:
|
||||
- action:
|
||||
- labels
|
||||
- nop
|
||||
initialize: testapi.testapiworkflow.TestInitFinal
|
||||
finalize: testapi.testapiworkflow.TestInitFinal
|
||||
merge: concat
|
||||
args:
|
||||
- [[positive, negative], false, True]
|
||||
- null
|
||||
schedule:
|
||||
schedule:
|
||||
cron: '* * * * * *'
|
||||
elements:
|
||||
- This is a test sentence. And another sentence to split.
|
||||
iterations: 1
|
||||
tasks:
|
||||
- action: segmentation
|
||||
segment:
|
||||
tasks:
|
||||
- action: segmentation
|
||||
- action: index
|
||||
get:
|
||||
tasks:
|
||||
- task: service
|
||||
url: http://127.0.0.1:8001/testget
|
||||
method: get
|
||||
params:
|
||||
text:
|
||||
post:
|
||||
tasks:
|
||||
- task: service
|
||||
url: http://127.0.0.1:8001/testpost
|
||||
params:
|
||||
|
||||
xml:
|
||||
tasks:
|
||||
- task: service
|
||||
url: http://127.0.0.1:8001/xml
|
||||
method: get
|
||||
batch: false
|
||||
extract: row
|
||||
params:
|
||||
text:
|
||||
"""
|
||||
|
||||
|
||||
class RequestHandler(BaseHTTPRequestHandler):
|
||||
"""
|
||||
Test HTTP handler.
|
||||
"""
|
||||
|
||||
def do_GET(self):
|
||||
"""
|
||||
GET request handler.
|
||||
"""
|
||||
|
||||
self.send_response(200)
|
||||
|
||||
if self.path.startswith("/xml"):
|
||||
response = "<row><text>test</text></row>".encode("utf-8")
|
||||
mime = "application/xml"
|
||||
else:
|
||||
response = '[{"text": "test"}]'.encode("utf-8")
|
||||
mime = "application/json"
|
||||
|
||||
self.send_header("content-type", mime)
|
||||
self.send_header("content-length", len(response))
|
||||
self.end_headers()
|
||||
|
||||
self.wfile.write(response)
|
||||
self.wfile.flush()
|
||||
|
||||
def do_POST(self):
|
||||
"""
|
||||
POST request handler.
|
||||
"""
|
||||
|
||||
length = int(self.headers["content-length"])
|
||||
data = json.loads(self.rfile.read(length))
|
||||
|
||||
response = json.dumps([[y for y in x.split(".") if y] for x in data]).encode("utf-8")
|
||||
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", "application/json")
|
||||
self.send_header("content-length", len(response))
|
||||
self.end_headers()
|
||||
|
||||
self.wfile.write(response)
|
||||
self.wfile.flush()
|
||||
|
||||
|
||||
class TestWorkflow(unittest.TestCase):
|
||||
"""
|
||||
API tests for workflows.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
|
||||
def start():
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testapi.yml")
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write(WORKFLOWS)
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestWorkflow.start()
|
||||
|
||||
cls.httpd = HTTPServer(("127.0.0.1", 8001), RequestHandler)
|
||||
|
||||
server = Thread(target=cls.httpd.serve_forever, daemon=True)
|
||||
server.start()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Shutdown mock http server.
|
||||
"""
|
||||
|
||||
cls.httpd.shutdown()
|
||||
|
||||
def testAPICleanup(self):
|
||||
"""
|
||||
Test API threadpool closed when __del__ called.
|
||||
"""
|
||||
|
||||
api = API({})
|
||||
api.pool = ThreadPool()
|
||||
|
||||
# pylint: disable=C2801
|
||||
api.__del__()
|
||||
|
||||
self.assertIsNone(api.pool)
|
||||
|
||||
def testServiceGet(self):
|
||||
"""
|
||||
Test workflow with ServiceTask GET via API
|
||||
"""
|
||||
|
||||
text = "This is a test sentence. And another sentence to split."
|
||||
results = self.client.post("workflow", json={"name": "get", "elements": [text]}).json()
|
||||
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(len(results[0]), 1)
|
||||
|
||||
def testServicePost(self):
|
||||
"""
|
||||
Test workflow with ServiceTask POST via API
|
||||
"""
|
||||
|
||||
text = "This is a test sentence. And another sentence to split."
|
||||
results = self.client.post("workflow", json={"name": "post", "elements": [text]}).json()
|
||||
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(len(results[0]), 2)
|
||||
|
||||
def testServiceXml(self):
|
||||
"""
|
||||
Test workflow with ServiceTask GET via API and XML response
|
||||
"""
|
||||
|
||||
text = "This is a test sentence. And another sentence to split."
|
||||
results = self.client.post("workflow", json={"name": "xml", "elements": [text]}).json()
|
||||
|
||||
self.assertEqual(len(results), 1)
|
||||
self.assertEqual(len(results[0]), 1)
|
||||
|
||||
def testWorkflowLabels(self):
|
||||
"""
|
||||
Test workflow with labels via API
|
||||
"""
|
||||
|
||||
text = "This is the best"
|
||||
|
||||
results = self.client.post("workflow", json={"name": "labels", "elements": [text]}).json()
|
||||
self.assertEqual(results[0][0], 0)
|
||||
|
||||
results = self.client.post("workflow", json={"name": "multiaction", "elements": [text]}).json()
|
||||
self.assertEqual(results[0], "['positive']. This is the best")
|
||||
|
||||
def testWorkflowSegment(self):
|
||||
"""
|
||||
Test workflow with segmentation via API
|
||||
"""
|
||||
|
||||
text = "This is a test sentence. And another sentence to split."
|
||||
|
||||
results = self.client.post("workflow", json={"name": "segment", "elements": [text]}).json()
|
||||
self.assertEqual(len(results), 2)
|
||||
|
||||
results = self.client.post("workflow", json={"name": "segment", "elements": [[0, text]]}).json()
|
||||
self.assertEqual(len(results), 2)
|
||||
|
||||
|
||||
class TestInitFinal:
|
||||
"""
|
||||
Class to test task initialize and finalize calls.
|
||||
"""
|
||||
|
||||
def __call__(self):
|
||||
pass
|
||||
@@ -0,0 +1,73 @@
|
||||
"""
|
||||
Authorization module tests
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import application
|
||||
|
||||
|
||||
class TestAuthorization(unittest.TestCase):
|
||||
"""
|
||||
API tests for token authorization.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"),
|
||||
"DEPENDENCIES": "txtai.api.Authorization",
|
||||
"TOKEN": hashlib.sha256("token".encode("utf-8")).hexdigest(),
|
||||
},
|
||||
)
|
||||
def start():
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testapi.yml")
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write("embeddings:\n")
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestAuthorization.start()
|
||||
|
||||
def testInvalid(self):
|
||||
"""
|
||||
Test invalid authorization
|
||||
"""
|
||||
|
||||
response = self.client.get("search?query=test")
|
||||
self.assertEqual(response.status_code, 401)
|
||||
|
||||
response = self.client.get("search?query=test", headers={"Authorization": "Bearer invalid"})
|
||||
self.assertEqual(response.status_code, 401)
|
||||
|
||||
def testValid(self):
|
||||
"""
|
||||
Test valid authorization
|
||||
"""
|
||||
|
||||
results = self.client.get("search?query=test", headers={"Authorization": "Bearer token"}).json()
|
||||
self.assertEqual(results, [])
|
||||
@@ -0,0 +1,261 @@
|
||||
"""
|
||||
Cluster API module tests
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
import urllib.parse
|
||||
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from threading import Thread
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import application
|
||||
|
||||
# Configuration for an embeddings cluster
|
||||
CLUSTER = """
|
||||
cluster:
|
||||
shards:
|
||||
- http://127.0.0.1:8002
|
||||
- http://127.0.0.1:8003
|
||||
"""
|
||||
|
||||
|
||||
class RequestHandler(BaseHTTPRequestHandler):
|
||||
"""
|
||||
Test HTTP handler.
|
||||
"""
|
||||
|
||||
def do_GET(self):
|
||||
"""
|
||||
GET request handler.
|
||||
"""
|
||||
|
||||
if self.path == "/count":
|
||||
response = 26
|
||||
elif self.path.startswith("/search?query=select"):
|
||||
if "group+by+id" in self.path:
|
||||
response = [{"count(*)": 26}]
|
||||
elif "group+by+text" in self.path:
|
||||
response = [{"count(*)": 12, "text": "This is a test"}, {"count(*)": 14, "text": "And another test"}]
|
||||
elif "group+by+txt" in self.path:
|
||||
response = [{"count(*)": 12, "txt": "This is a test"}, {"count(*)": 14, "txt": "And another test"}]
|
||||
else:
|
||||
if self.server.server_port == 8002:
|
||||
response = [{"count(*)": 12, "min(indexid)": 0, "max(indexid)": 11, "avg(indexid)": 6.3}]
|
||||
else:
|
||||
response = [{"count(*)": 16, "min(indexid)": 2, "max(indexid)": 14, "avg(indexid)": 6.7}]
|
||||
elif self.path.startswith("/search"):
|
||||
response = [{"id": 4, "score": 0.40}]
|
||||
else:
|
||||
response = {"result": "ok"}
|
||||
|
||||
# Convert response to string
|
||||
response = json.dumps(response).encode("utf-8")
|
||||
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", "application/json")
|
||||
self.send_header("content-length", len(response))
|
||||
self.end_headers()
|
||||
|
||||
self.wfile.write(response)
|
||||
self.wfile.flush()
|
||||
|
||||
def do_POST(self):
|
||||
"""
|
||||
POST request handler.
|
||||
"""
|
||||
|
||||
if self.path.startswith("/batchsearch"):
|
||||
response = [[{"id": 4, "score": 0.40}], [{"id": 1, "score": 0.40}]]
|
||||
elif self.path.startswith("/delete"):
|
||||
if self.server.server_port == 8002:
|
||||
response = [0]
|
||||
else:
|
||||
response = []
|
||||
else:
|
||||
response = {"result": "ok"}
|
||||
|
||||
response = json.dumps(response).encode("utf-8")
|
||||
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", "application/json")
|
||||
self.send_header("content-length", len(response))
|
||||
self.end_headers()
|
||||
|
||||
self.wfile.write(response)
|
||||
self.wfile.flush()
|
||||
|
||||
|
||||
@unittest.skipIf(os.name == "nt", "TestCluster skipped on Windows")
|
||||
class TestCluster(unittest.TestCase):
|
||||
"""
|
||||
API tests for embeddings clusters
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
|
||||
def start():
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testapi.yml")
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write(CLUSTER)
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestCluster.start()
|
||||
|
||||
cls.httpd1 = HTTPServer(("127.0.0.1", 8002), RequestHandler)
|
||||
|
||||
server1 = Thread(target=cls.httpd1.serve_forever, daemon=True)
|
||||
server1.start()
|
||||
|
||||
cls.httpd2 = HTTPServer(("127.0.0.1", 8003), RequestHandler)
|
||||
|
||||
server2 = Thread(target=cls.httpd2.serve_forever, daemon=True)
|
||||
server2.start()
|
||||
|
||||
# Index data
|
||||
cls.client.post("add", json=[{"id": 0, "text": "test"}])
|
||||
cls.client.get("index")
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Shutdown mock http server.
|
||||
"""
|
||||
|
||||
cls.httpd1.shutdown()
|
||||
cls.httpd2.shutdown()
|
||||
|
||||
def testCount(self):
|
||||
"""
|
||||
Test cluster count
|
||||
"""
|
||||
|
||||
self.assertEqual(self.client.get("count").json(), 52)
|
||||
|
||||
def testDelete(self):
|
||||
"""
|
||||
Test cluster delete
|
||||
"""
|
||||
|
||||
self.assertEqual(self.client.post("delete", json=[0]).json(), [0])
|
||||
|
||||
def testDeleteString(self):
|
||||
"""
|
||||
Test cluster delete with string id
|
||||
"""
|
||||
|
||||
self.assertEqual(self.client.post("delete", json=["0"]).json(), [0])
|
||||
|
||||
def testIds(self):
|
||||
"""
|
||||
Test id configurations
|
||||
"""
|
||||
|
||||
# String ids
|
||||
self.client.post("add", json=[{"id": "0", "text": "test"}])
|
||||
self.assertEqual(self.client.get("index").status_code, 200)
|
||||
|
||||
# Auto ids
|
||||
self.client.post("add", json=[{"text": "test"}])
|
||||
self.assertEqual(self.client.get("index").status_code, 200)
|
||||
|
||||
def testReindex(self):
|
||||
"""
|
||||
Test cluster reindex
|
||||
"""
|
||||
|
||||
self.assertEqual(self.client.post("reindex", json={"config": {"path": "sentence-transformers/nli-mpnet-base-v2"}}).status_code, 200)
|
||||
|
||||
def testSearch(self):
|
||||
"""
|
||||
Test cluster search
|
||||
"""
|
||||
|
||||
# Encode parameters
|
||||
params = json.dumps({"x": 1})
|
||||
|
||||
query = urllib.parse.quote("feel good story")
|
||||
uid = self.client.get(f"search?query={query}&limit=1&weights=0.5&index=default¶meters={params}&graph=False").json()[0]["id"]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testSearchBatch(self):
|
||||
"""
|
||||
Test cluster batch search
|
||||
"""
|
||||
|
||||
results = self.client.post(
|
||||
"batchsearch",
|
||||
json={
|
||||
"queries": ["feel good story", "climate change"],
|
||||
"limit": 1,
|
||||
"weights": 0.5,
|
||||
"index": "default",
|
||||
"parameters": [{"x": 1}, {"x": 2}],
|
||||
"graph": False,
|
||||
},
|
||||
).json()
|
||||
|
||||
uids = [result[0]["id"] for result in results]
|
||||
self.assertEqual(uids, [4, 1])
|
||||
|
||||
def testSQL(self):
|
||||
"""
|
||||
Test cluster SQL statement
|
||||
"""
|
||||
|
||||
query = urllib.parse.quote("select count(*), min(indexid), max(indexid), avg(indexid) from txtai where text='This is a test'")
|
||||
self.assertEqual(
|
||||
self.client.get(f"search?query={query}").json(), [{"count(*)": 28, "min(indexid)": 0, "max(indexid)": 14, "avg(indexid)": 6.5}]
|
||||
)
|
||||
|
||||
query = urllib.parse.quote("select count(*), text txt from txtai group by txt order by count(*) desc")
|
||||
self.assertEqual(
|
||||
self.client.get(f"search?query={query}").json(),
|
||||
[{"count(*)": 28, "txt": "And another test"}, {"count(*)": 24, "txt": "This is a test"}],
|
||||
)
|
||||
|
||||
query = urllib.parse.quote("select count(*), text from txtai group by text order by count(*) asc")
|
||||
self.assertEqual(
|
||||
self.client.get(f"search?query={query}").json(),
|
||||
[{"count(*)": 24, "text": "This is a test"}, {"count(*)": 28, "text": "And another test"}],
|
||||
)
|
||||
|
||||
query = urllib.parse.quote("select count(*) from txtai group by id order by count(*)")
|
||||
self.assertEqual(self.client.get(f"search?query={query}").json(), [{"count(*)": 52}])
|
||||
|
||||
def testUpsert(self):
|
||||
"""
|
||||
Test cluster upsert
|
||||
"""
|
||||
|
||||
# Update data
|
||||
self.client.post("add", json=[{"id": 4, "text": "Feel good story: baby panda born"}])
|
||||
self.client.get("upsert")
|
||||
|
||||
# Search for best match
|
||||
query = "feel good story"
|
||||
uid = self.client.get(f"search?query={query}&limit=1").json()[0]["id"]
|
||||
|
||||
self.assertEqual(uid, 4)
|
||||
@@ -0,0 +1,163 @@
|
||||
"""
|
||||
Encoding module tests
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
import urllib.parse
|
||||
|
||||
from io import BytesIO
|
||||
from unittest.mock import patch
|
||||
|
||||
import msgpack
|
||||
import numpy as np
|
||||
import PIL
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import application
|
||||
from txtai.api.responses import JSONEncoder
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
# Configuration for image storage
|
||||
INDEX = """
|
||||
# Allow indexing of documents
|
||||
writable: %s
|
||||
|
||||
embeddings:
|
||||
defaults: False
|
||||
content: True
|
||||
objects: %s
|
||||
"""
|
||||
|
||||
|
||||
class TestEncoding(unittest.TestCase):
|
||||
"""
|
||||
API tests for response encoding
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
|
||||
def start(yaml):
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
|
||||
Args:
|
||||
yaml: input configuration
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testapi.yml")
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write(yaml)
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestEncoding.start(INDEX % ("True", "image"))
|
||||
|
||||
def testImages(self):
|
||||
"""
|
||||
Test image encoding
|
||||
"""
|
||||
|
||||
with open(Utils.PATH + "/books.jpg", "rb") as f:
|
||||
self.client.post("addimage", data={"uid": 0}, files={"data": f})
|
||||
self.client.get("index")
|
||||
|
||||
query = urllib.parse.quote_plus("select id, object from txtai limit 1")
|
||||
results = self.client.get(f"search?query={query}").json()
|
||||
|
||||
# Test reading image
|
||||
self.assertIsInstance(PIL.Image.open(BytesIO(base64.b64decode(results[0]["object"]))), PIL.Image.Image)
|
||||
|
||||
def testInvalidInputs(self):
|
||||
"""
|
||||
Test invalid parameter inputs
|
||||
"""
|
||||
|
||||
response = self.client.post("addimage", data={"uid": [0, 1]}, files={"data": b"123"})
|
||||
self.assertEqual(response.status_code, 422)
|
||||
|
||||
response = self.client.post("addobject", data={"uid": [0, 1]}, files={"data": b"123"})
|
||||
self.assertEqual(response.status_code, 422)
|
||||
|
||||
def testInvalidJSON(self):
|
||||
"""
|
||||
Test that invalid JSON raises an exception
|
||||
"""
|
||||
|
||||
with self.assertRaises(TypeError):
|
||||
JSONEncoder().encode(np.random.rand(1, 1))
|
||||
|
||||
def testMessagePack(self):
|
||||
"""
|
||||
Test message pack encoding
|
||||
"""
|
||||
|
||||
# Validate binary encoding
|
||||
results = self.client.get("count", headers={"Accept": "application/msgpack"}).content
|
||||
self.assertEqual(results, b"\x01")
|
||||
|
||||
# Validate query result
|
||||
query = urllib.parse.quote_plus("select id, object from txtai limit 1")
|
||||
results = self.client.get(f"search?query={query}", headers={"Accept": "application/msgpack"}).content
|
||||
results = msgpack.unpackb(results)
|
||||
|
||||
# Test reading image
|
||||
self.assertIsInstance(PIL.Image.open(BytesIO(results[0]["object"])), PIL.Image.Image)
|
||||
|
||||
def testObjects(self):
|
||||
"""
|
||||
Test object encoding
|
||||
"""
|
||||
|
||||
# Recreate model with standard object encoding
|
||||
self.client = TestEncoding.start(INDEX % ("True", "True"))
|
||||
|
||||
# Test various formats
|
||||
self.client.post("addobject", data={"uid": "id0"}, files={"data": b"1234"})
|
||||
self.client.post("addobject", files={"data": b"ABC"})
|
||||
self.client.post("addobject", data={"uid": "id1", "field": "object"}, files={"data": b"A1234"})
|
||||
self.client.get("index")
|
||||
|
||||
query = urllib.parse.quote_plus("select id, object from txtai where id = 'id0' limit 1")
|
||||
results = self.client.get(f"search?query={query}").json()
|
||||
self.assertEqual(base64.b64decode(results[0]["object"]), b"1234")
|
||||
|
||||
# Test with messagepack encoding
|
||||
results = self.client.get(f"search?query={query}", headers={"Accept": "application/msgpack"}).content
|
||||
results = msgpack.unpackb(results)
|
||||
self.assertEqual(results[0]["object"], b"1234")
|
||||
|
||||
count = self.client.get("count").json()
|
||||
self.assertEqual(count, 3)
|
||||
|
||||
def testReadOnly(self):
|
||||
"""
|
||||
Test read only indexes
|
||||
"""
|
||||
|
||||
# Recreate model with standard object encoding
|
||||
self.client = TestEncoding.start(INDEX % ("False", "True"))
|
||||
|
||||
# Test errors raised for write operations
|
||||
with open(Utils.PATH + "/books.jpg", "rb") as f:
|
||||
response = self.client.post("addimage", data={"uid": 0}, files={"data": f})
|
||||
self.assertEqual(response.status_code, 403)
|
||||
|
||||
self.assertEqual(self.client.post("addobject", data={"uid": 0}, files={"data": b"1234"}).status_code, 403)
|
||||
@@ -0,0 +1,117 @@
|
||||
"""
|
||||
Extension module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi import APIRouter
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import application, Extension
|
||||
from txtai.pipeline import Pipeline
|
||||
|
||||
# Example pipeline extension
|
||||
PIPELINES = """
|
||||
testapi.testextension.SamplePipeline:
|
||||
"""
|
||||
|
||||
|
||||
class SampleRouter:
|
||||
"""
|
||||
Sample API router.
|
||||
"""
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@staticmethod
|
||||
@router.get("/sample")
|
||||
def sample(text: str):
|
||||
"""
|
||||
Calls sample pipeline.
|
||||
|
||||
Args:
|
||||
text: input text
|
||||
|
||||
Returns:
|
||||
formatted text
|
||||
"""
|
||||
|
||||
return application.get().pipeline("testapi.testextension.SamplePipeline", (text,))
|
||||
|
||||
|
||||
class SampleExtension(Extension):
|
||||
"""
|
||||
Sample API extension.
|
||||
"""
|
||||
|
||||
def __call__(self, app):
|
||||
app.include_router(SampleRouter().router)
|
||||
|
||||
|
||||
class SamplePipeline(Pipeline):
|
||||
"""
|
||||
Sample pipeline.
|
||||
"""
|
||||
|
||||
def __call__(self, text):
|
||||
return text.lower()
|
||||
|
||||
|
||||
class TestExtension(unittest.TestCase):
|
||||
"""
|
||||
API tests for extensions.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"),
|
||||
"API_CLASS": "txtai.api.API",
|
||||
"EXTENSIONS": "testapi.testextension.SampleExtension",
|
||||
},
|
||||
)
|
||||
def start():
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testapi.yml")
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write(PIPELINES)
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestExtension.start()
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test an empty extension
|
||||
"""
|
||||
|
||||
extension = Extension()
|
||||
self.assertIsNone(extension(None))
|
||||
|
||||
def testExtension(self):
|
||||
"""
|
||||
Test a pipeline extension
|
||||
"""
|
||||
|
||||
text = self.client.get("sample?text=Test%20String").json()
|
||||
self.assertEqual(text, "test string")
|
||||
@@ -0,0 +1,59 @@
|
||||
"""
|
||||
Agent API module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import application
|
||||
|
||||
# Configuration for agents
|
||||
MCP = """
|
||||
mcp: True
|
||||
"""
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestMCP(unittest.TestCase):
|
||||
"""
|
||||
API tests for model context protocol (MCP)
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testapi.yml"), "API_CLASS": "txtai.api.API"})
|
||||
def start():
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testapi.yml")
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write(MCP)
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestMCP.start()
|
||||
|
||||
def testMCP(self):
|
||||
"""
|
||||
Test that application a /mcp route
|
||||
"""
|
||||
|
||||
self.assertTrue(any(route.path == "/mcp" for route in self.client.app.routes))
|
||||
@@ -0,0 +1,201 @@
|
||||
"""
|
||||
OpenAI API module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from txtai.api import application
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
# API Configuration
|
||||
CONFIG = """
|
||||
# Enable OpenAI-compatible API
|
||||
openai: True
|
||||
|
||||
# Allow indexing of documents
|
||||
writable: True
|
||||
|
||||
# Agent configuration
|
||||
agent:
|
||||
hello:
|
||||
max_iterations: 1
|
||||
|
||||
# Embeddings settings
|
||||
embeddings:
|
||||
path: sentence-transformers/nli-mpnet-base-v2
|
||||
content: True
|
||||
|
||||
# LLM configuration
|
||||
llm:
|
||||
path: hf-internal-testing/tiny-random-LlamaForCausalLM
|
||||
|
||||
# Text segmentation
|
||||
segmentation:
|
||||
|
||||
# Text to speech
|
||||
texttospeech:
|
||||
|
||||
# Transcription
|
||||
transcription:
|
||||
|
||||
# Workflow
|
||||
workflow:
|
||||
echo:
|
||||
tasks:
|
||||
- task: console
|
||||
"""
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestOpenAI(unittest.TestCase):
|
||||
"""
|
||||
Tests for OpenAI-compatible API endpoint for txtai.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
@patch.dict(os.environ, {"CONFIG": os.path.join(tempfile.gettempdir(), "testopenai.yml"), "API_CLASS": "txtai.api.API"})
|
||||
def start():
|
||||
"""
|
||||
Starts a mock FastAPI client.
|
||||
"""
|
||||
|
||||
config = os.path.join(tempfile.gettempdir(), "testopenai.yml")
|
||||
|
||||
with open(config, "w", encoding="utf-8") as output:
|
||||
output.write(CONFIG)
|
||||
|
||||
# Create new application and set on client
|
||||
application.app = application.create()
|
||||
client = TestClient(application.app)
|
||||
application.start()
|
||||
|
||||
# Patch LLM to generate answer
|
||||
agent = application.get().agents["hello"]
|
||||
agent.process.model.llm = lambda *args, **kwargs: 'Action:\n{"name": "final_answer", "arguments": "Hi"}'
|
||||
|
||||
return client
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create API client on creation of class.
|
||||
"""
|
||||
|
||||
cls.client = TestOpenAI.start()
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Index data
|
||||
cls.client.post("add", json=[{"id": x, "text": row} for x, row in enumerate(cls.data)])
|
||||
cls.client.get("index")
|
||||
|
||||
def testChatAgent(self):
|
||||
"""
|
||||
Test a chat completion with an agent
|
||||
"""
|
||||
|
||||
response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "hello"}).json()
|
||||
|
||||
self.assertEqual(response["choices"][0]["message"]["content"], "Hi")
|
||||
|
||||
def testChatLLM(self):
|
||||
"""
|
||||
Test a chat completion with a LLM
|
||||
"""
|
||||
|
||||
response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "llm"}).json()
|
||||
|
||||
self.assertIsNotNone(response["choices"][0]["message"]["content"])
|
||||
|
||||
def testChatPipeline(self):
|
||||
"""
|
||||
Test a chat completion with a pipeline
|
||||
"""
|
||||
|
||||
response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "segmentation"}).json()
|
||||
|
||||
self.assertEqual(response["choices"][0]["message"]["content"], "Hello")
|
||||
|
||||
def testChatSearch(self):
|
||||
"""
|
||||
Test a chat completion with an embeddings search
|
||||
"""
|
||||
|
||||
response = self.client.post(
|
||||
"/v1/chat/completions", json={"messages": [{"role": "user", "content": "feel good story"}], "model": "embeddings"}
|
||||
).json()
|
||||
|
||||
self.assertEqual(response["choices"][0]["message"]["content"], self.data[4])
|
||||
|
||||
def testChatStream(self):
|
||||
"""
|
||||
Test a chat completion with a LLM
|
||||
"""
|
||||
|
||||
response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "llm", "stream": True})
|
||||
|
||||
self.assertGreater(len(response.text.split("\n\n")), 0)
|
||||
|
||||
def testChatWorkflow(self):
|
||||
"""
|
||||
Test a chat completion with a workflow
|
||||
"""
|
||||
|
||||
response = self.client.post("/v1/chat/completions", json={"messages": [{"role": "user", "content": "Hello"}], "model": "echo"}).json()
|
||||
|
||||
self.assertEqual(response["choices"][0]["message"]["content"], "Hello")
|
||||
|
||||
def testEmbeddings(self):
|
||||
"""
|
||||
Test generating embeddings vectors
|
||||
"""
|
||||
|
||||
response = self.client.post("/v1/embeddings", json={"input": "text to embed", "model": "nli-mpnet-base-v2"}).json()
|
||||
|
||||
self.assertEqual(len(response["data"][0]["embedding"]), 768)
|
||||
|
||||
def testSpeech(self):
|
||||
"""
|
||||
Test generating speech for input text
|
||||
"""
|
||||
|
||||
response = self.client.post(
|
||||
"/v1/audio/speech", json={"model": "tts", "input": "text to speak", "voice": "default", "response_format": "wav"}
|
||||
).content
|
||||
|
||||
self.assertTrue(response[0:4] == b"RIFF")
|
||||
|
||||
def testTranscribe(self):
|
||||
"""
|
||||
Test audio to text transcription
|
||||
"""
|
||||
|
||||
path = Utils.PATH + "/Make_huge_profits.wav"
|
||||
with open(path, "rb") as f:
|
||||
text = self.client.post("/v1/audio/transcriptions", files={"file": f}).json()["text"]
|
||||
self.assertEqual(text, "Make huge profits without working make up to one hundred thousand dollars a day")
|
||||
|
||||
def testTranslate(self):
|
||||
"""
|
||||
Test audio translation
|
||||
"""
|
||||
|
||||
path = Utils.PATH + "/Make_huge_profits.wav"
|
||||
with open(path, "rb") as f:
|
||||
text = self.client.post("/v1/audio/translations", files={"file": f}).json()["text"]
|
||||
self.assertEqual(text, "Make huge profits without working make up to one hundred thousand dollars a day")
|
||||
@@ -0,0 +1,91 @@
|
||||
"""
|
||||
Application module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
import types
|
||||
|
||||
from txtai.app import Application
|
||||
from txtai.pipeline import Pipeline
|
||||
|
||||
|
||||
class TestApp(unittest.TestCase):
|
||||
"""
|
||||
Application tests.
|
||||
"""
|
||||
|
||||
def testConfig(self):
|
||||
"""
|
||||
Test a file not found config exception
|
||||
"""
|
||||
|
||||
with self.assertRaises(FileNotFoundError):
|
||||
Application.read("No file here")
|
||||
|
||||
def testParameter(self):
|
||||
"""
|
||||
Test resolving application parameter
|
||||
"""
|
||||
|
||||
app = Application(
|
||||
"""
|
||||
testapp.TestPipeline:
|
||||
application:
|
||||
"""
|
||||
)
|
||||
|
||||
# Check that application instance is not None
|
||||
self.assertIsNotNone(app.pipelines["testapp.TestPipeline"].application)
|
||||
|
||||
def testStream(self):
|
||||
"""
|
||||
Test workflow streams
|
||||
"""
|
||||
|
||||
app = Application(
|
||||
"""
|
||||
workflow:
|
||||
stream:
|
||||
stream:
|
||||
action: testapp.TestStream
|
||||
tasks:
|
||||
- nop
|
||||
batchstream:
|
||||
stream:
|
||||
action: testapp.TestStream
|
||||
batch: True
|
||||
tasks:
|
||||
- nop
|
||||
"""
|
||||
)
|
||||
|
||||
def generator():
|
||||
yield 10
|
||||
|
||||
# Test single stream
|
||||
self.assertEqual(list(app.workflow("stream", [10])), list(range(10)))
|
||||
|
||||
# Test batch stream
|
||||
self.assertEqual(list(app.workflow("batchstream", generator())), list(range(10)))
|
||||
|
||||
|
||||
class TestPipeline(Pipeline):
|
||||
"""
|
||||
Test pipeline with an application parameter.
|
||||
"""
|
||||
|
||||
def __init__(self, application):
|
||||
self.application = application
|
||||
|
||||
|
||||
class TestStream:
|
||||
"""
|
||||
Test workflow stream
|
||||
"""
|
||||
|
||||
def __call__(self, arg):
|
||||
if isinstance(arg, types.GeneratorType):
|
||||
for x in arg:
|
||||
yield from range(int(x))
|
||||
else:
|
||||
yield from range(int(arg))
|
||||
@@ -0,0 +1,116 @@
|
||||
"""
|
||||
Compress module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tarfile
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from zipfile import ZipFile, ZIP_DEFLATED
|
||||
|
||||
from txtai.archive import ArchiveFactory, Compress
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestArchive(unittest.TestCase):
|
||||
"""
|
||||
Archive tests.
|
||||
"""
|
||||
|
||||
def testDirectory(self):
|
||||
"""
|
||||
Test directory included in compressed files
|
||||
"""
|
||||
|
||||
for extension in ["tar", "zip"]:
|
||||
# Create archive instance
|
||||
archive = ArchiveFactory.create()
|
||||
|
||||
# Create subdirectory in archive working path
|
||||
path = os.path.join(archive.path(), "dir")
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
# Create file in archive working path
|
||||
with open(os.path.join(path, "test"), "w", encoding="utf-8") as f:
|
||||
f.write("test")
|
||||
|
||||
# Save archive
|
||||
path = os.path.join(tempfile.gettempdir(), f"subdir.{extension}")
|
||||
archive.save(path)
|
||||
|
||||
# Extract files from archive
|
||||
archive = ArchiveFactory.create()
|
||||
archive.load(path)
|
||||
|
||||
# Check if file properly extracted
|
||||
path = os.path.join(archive.path(), "dir", "test")
|
||||
self.assertTrue(os.path.exists(path))
|
||||
|
||||
def testInvalidTarLink(self):
|
||||
"""
|
||||
Test invalid tar file with symlinks
|
||||
"""
|
||||
|
||||
symlink = os.path.join(tempfile.gettempdir(), "link")
|
||||
|
||||
# Remove symlink if it already exists
|
||||
try:
|
||||
os.remove(symlink)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
# Create symlink and add to TAR file
|
||||
os.symlink(os.path.join(tempfile.gettempdir(), "noexist"), symlink)
|
||||
|
||||
path = os.path.join(tempfile.gettempdir(), "badtarlink")
|
||||
with tarfile.open(path, "w") as tar:
|
||||
tar.add(symlink, arcname="l")
|
||||
|
||||
archive = ArchiveFactory.create()
|
||||
|
||||
# Validate error is thrown for file
|
||||
with self.assertRaises(IOError):
|
||||
archive.load(path, "tar")
|
||||
|
||||
def testInvalidTarPath(self):
|
||||
"""
|
||||
Test invalid tar file with a path outside of base directory
|
||||
"""
|
||||
|
||||
path = os.path.join(tempfile.gettempdir(), "badtarpath")
|
||||
with tarfile.open(path, "w") as tar:
|
||||
tar.add(Utils.PATH, arcname="..")
|
||||
|
||||
archive = ArchiveFactory.create()
|
||||
|
||||
# Validate error is thrown for file
|
||||
with self.assertRaises(IOError):
|
||||
archive.load(path, "tar")
|
||||
|
||||
def testInvalidZipPath(self):
|
||||
"""
|
||||
Test invalid zip file with a path outside of base directory
|
||||
"""
|
||||
|
||||
path = os.path.join(tempfile.gettempdir(), "badzippath")
|
||||
with ZipFile(path, "w", ZIP_DEFLATED) as zfile:
|
||||
zfile.write(Utils.PATH + "/article.pdf", arcname="../article.pdf")
|
||||
|
||||
archive = ArchiveFactory.create()
|
||||
|
||||
# Validate error is thrown for file
|
||||
with self.assertRaises(IOError):
|
||||
archive.load(path, "zip")
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
compress = Compress()
|
||||
|
||||
self.assertRaises(NotImplementedError, compress.pack, None, None)
|
||||
self.assertRaises(NotImplementedError, compress.unpack, None, None)
|
||||
@@ -0,0 +1,181 @@
|
||||
"""
|
||||
Cloud module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
from txtai.cloud import Cloud
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
|
||||
class TestCloud(unittest.TestCase):
|
||||
"""
|
||||
Cloud tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"format": "json", "path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def testCustom(self):
|
||||
"""
|
||||
Test custom provider
|
||||
"""
|
||||
|
||||
# pylint: disable=E1120
|
||||
self.runHub("txtai.cloud.HuggingFaceHub")
|
||||
|
||||
def testHub(self):
|
||||
"""
|
||||
Test huggingface-hub integration
|
||||
"""
|
||||
|
||||
# pylint: disable=E1120
|
||||
self.runHub("huggingface-hub")
|
||||
|
||||
def testInvalidProvider(self):
|
||||
"""
|
||||
Test invalid provider identifier
|
||||
"""
|
||||
|
||||
# Test invalid external provider
|
||||
with self.assertRaises(ImportError):
|
||||
embeddings = Embeddings()
|
||||
embeddings.load(provider="ProviderNoExist", container="Invalid")
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
cloud = Cloud({})
|
||||
|
||||
self.assertRaises(NotImplementedError, cloud.exists, None)
|
||||
self.assertRaises(NotImplementedError, cloud.load, None)
|
||||
self.assertRaises(NotImplementedError, cloud.save, None)
|
||||
|
||||
def testObjectStorage(self):
|
||||
"""
|
||||
Test object storage integration
|
||||
"""
|
||||
|
||||
# Run tests with uncompressed and compressed index
|
||||
for path in ["cloud.object", "cloud.object.tar.gz"]:
|
||||
self.runTests(path, {"provider": "local", "container": f"cloud.{time.time()}", "key": tempfile.gettempdir()})
|
||||
|
||||
@patch("huggingface_hub.hf_hub_download")
|
||||
@patch("huggingface_hub.get_hf_file_metadata")
|
||||
@patch("huggingface_hub.upload_file")
|
||||
@patch("huggingface_hub.create_repo")
|
||||
def runHub(self, provider, create, upload, metadata, download):
|
||||
"""
|
||||
Run huggingface-hub tests. This method mocks write operations since a token won't be available.
|
||||
"""
|
||||
|
||||
def filemeta(url, token):
|
||||
return (url, token) if "Invalid" not in url else None
|
||||
|
||||
def filedownload(**kwargs):
|
||||
if "Invalid" in kwargs["repo_id"]:
|
||||
raise FileNotFoundError
|
||||
|
||||
# Check for .gitattributes file
|
||||
if kwargs["filename"] == ".gitattributes":
|
||||
return attributes
|
||||
|
||||
# Check for cloud index path
|
||||
if any(kwargs["filename"] == x for x in paths):
|
||||
return index
|
||||
|
||||
# Use original method for all other requests
|
||||
return hf_hub_download(**kwargs)
|
||||
|
||||
# Patch write methods since token will not be available
|
||||
create.return_value = None
|
||||
upload.return_value = None
|
||||
metadata.side_effect = filemeta
|
||||
download.side_effect = filedownload
|
||||
|
||||
# Create dummy index
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"cloud.{provider}.tar.gz")
|
||||
self.embeddings.save(index)
|
||||
|
||||
# Initialize attributes file
|
||||
# pylint: disable=R1732
|
||||
with tempfile.NamedTemporaryFile(mode="w", delete=False) as tmp:
|
||||
tmp.write("*.bin filter=lfs diff=lfs merge=lfs -text\n")
|
||||
attributes = tmp.name
|
||||
|
||||
# Run tests with uncompressed and compressed index
|
||||
paths = [f"cloud.{provider}", f"cloud.{provider}.tar.gz"]
|
||||
for path in paths:
|
||||
self.runTests(path, {"provider": provider, "container": "neuml/txtai-intro"})
|
||||
|
||||
def runTests(self, path, cloud):
|
||||
"""
|
||||
Runs a series of cloud sync tests.
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), path)
|
||||
|
||||
# Test exists handles missing cloud storage object
|
||||
invalid = cloud.copy()
|
||||
invalid["container"] = "InvalidPathToTest"
|
||||
self.assertFalse(self.embeddings.exists(index, invalid))
|
||||
|
||||
# Test exception raised when trying to load index and doesn't exist in cloud storage
|
||||
# pylint: disable=W0719
|
||||
with self.assertRaises(Exception):
|
||||
self.embeddings.load(index, invalid)
|
||||
|
||||
# Save index
|
||||
self.embeddings.save(index, cloud)
|
||||
|
||||
# Test object exists in cloud storage
|
||||
self.assertTrue(self.embeddings.exists(index, cloud))
|
||||
|
||||
# Test object exists locally
|
||||
self.assertTrue(self.embeddings.exists(index))
|
||||
|
||||
# Test index can be reloaded
|
||||
self.embeddings.load(index, cloud)
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
@@ -0,0 +1,175 @@
|
||||
"""
|
||||
Console module tests
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from txtai.console import Console
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
APPLICATION = """
|
||||
path: %s
|
||||
|
||||
workflow:
|
||||
test:
|
||||
tasks:
|
||||
- task: console
|
||||
"""
|
||||
|
||||
|
||||
class TestConsole(unittest.TestCase):
|
||||
"""
|
||||
Console tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
|
||||
|
||||
# Create an index for the list of text
|
||||
cls.embeddings.index([(uid, text, None) for uid, text in enumerate(cls.data)])
|
||||
|
||||
# Create app paths
|
||||
cls.apppath = os.path.join(tempfile.gettempdir(), "console.yml")
|
||||
cls.embedpath = os.path.join(tempfile.gettempdir(), "embeddings.console")
|
||||
|
||||
# Create app.yml
|
||||
with open(cls.apppath, "w", encoding="utf-8") as out:
|
||||
out.write(APPLICATION % cls.embedpath)
|
||||
|
||||
# Save index as uncompressed and compressed
|
||||
cls.embeddings.save(cls.embedpath)
|
||||
cls.embeddings.save(f"{cls.embedpath}.tar.gz")
|
||||
|
||||
# Create console
|
||||
cls.console = Console(cls.embedpath)
|
||||
|
||||
def testApplication(self):
|
||||
"""
|
||||
Test application
|
||||
"""
|
||||
|
||||
self.assertNotIn("Traceback", self.command(f".load {self.apppath}"))
|
||||
self.assertIn("1", self.command(".limit 1"))
|
||||
self.assertIn("Maine man wins", self.command("feel good story"))
|
||||
|
||||
def testConfig(self):
|
||||
"""
|
||||
Test .config command
|
||||
"""
|
||||
|
||||
self.assertIn("tasks", self.command(".config"))
|
||||
|
||||
def testEmbeddings(self):
|
||||
"""
|
||||
Test embeddings index
|
||||
"""
|
||||
|
||||
self.assertNotIn("Traceback", self.command(f".load {self.embedpath}.tar.gz"))
|
||||
self.assertNotIn("Traceback", self.command(f".load {self.embedpath}"))
|
||||
self.assertIn("1", self.command(".limit 1"))
|
||||
self.assertIn("Maine man wins", self.command("feel good story"))
|
||||
|
||||
def testEmbeddingsNoDatabase(self):
|
||||
"""
|
||||
Test embeddings with no database/content
|
||||
"""
|
||||
|
||||
console = Console()
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Set embeddings on console
|
||||
console.app = embeddings
|
||||
self.assertIn("4", self.command("feel good story", console))
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test empty console instance
|
||||
"""
|
||||
|
||||
console = Console()
|
||||
self.assertIn("AttributeError", self.command("search", console))
|
||||
|
||||
def testHighlight(self):
|
||||
"""
|
||||
Test .highlight command
|
||||
"""
|
||||
|
||||
self.assertIn("highlight", self.command(".highlight"))
|
||||
self.assertIn("wins", self.command("feel good story"))
|
||||
self.assertIn("Taiwan", self.command("asia"))
|
||||
|
||||
def testPreloop(self):
|
||||
"""
|
||||
Test preloop
|
||||
"""
|
||||
|
||||
self.assertIn("txtai console", self.preloop())
|
||||
|
||||
def testWorkflow(self):
|
||||
"""
|
||||
Test .workflow command
|
||||
"""
|
||||
|
||||
self.command(f".load {self.apppath}")
|
||||
self.assertIn("echo", self.command(".workflow test echo"))
|
||||
|
||||
def command(self, command, console=None):
|
||||
"""
|
||||
Runs a console command.
|
||||
|
||||
Args:
|
||||
command: command to run
|
||||
console: console instance, defaults to self.console
|
||||
|
||||
Returns:
|
||||
command output
|
||||
"""
|
||||
|
||||
# Run info
|
||||
output = io.StringIO()
|
||||
with contextlib.redirect_stdout(output):
|
||||
if not console:
|
||||
console = self.console
|
||||
|
||||
console.onecmd(command)
|
||||
|
||||
return output.getvalue()
|
||||
|
||||
def preloop(self):
|
||||
"""
|
||||
Runs console.preloop and redirects stdout.
|
||||
|
||||
Returns:
|
||||
preloop output
|
||||
"""
|
||||
|
||||
# Run info
|
||||
output = io.StringIO()
|
||||
with contextlib.redirect_stdout(output):
|
||||
self.console.preloop()
|
||||
|
||||
return output.getvalue()
|
||||
@@ -0,0 +1,71 @@
|
||||
"""
|
||||
Client module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import tempfile
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
from .testrdbms import Common
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestClient(Common.TestRDBMS):
|
||||
"""
|
||||
Embeddings with content stored in a client RDBMS.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Content backend
|
||||
cls.backend = None
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def setUp(self):
|
||||
"""
|
||||
Set unique database path for each test.
|
||||
"""
|
||||
|
||||
# Generate unique database path and set on embeddings
|
||||
path = os.path.join(tempfile.gettempdir(), f"{int(time.time() * 1000)}.sqlite")
|
||||
self.backend = f"sqlite:///{path}"
|
||||
|
||||
self.embeddings.config["content"] = self.backend
|
||||
|
||||
def testSchema(self):
|
||||
"""
|
||||
Test database creation with a specified schema
|
||||
"""
|
||||
|
||||
# Default sequence id
|
||||
embeddings = Embeddings(path="sentence-transformers/nli-mpnet-base-v2", content=self.backend, schema="txtai")
|
||||
embeddings.index(self.data)
|
||||
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
@@ -0,0 +1,29 @@
|
||||
"""
|
||||
Custom database tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.database import DatabaseFactory
|
||||
|
||||
|
||||
class TestCustom(unittest.TestCase):
|
||||
"""
|
||||
Custom database backend tests.
|
||||
"""
|
||||
|
||||
def testCustomBackend(self):
|
||||
"""
|
||||
Test resolving a custom backend
|
||||
"""
|
||||
|
||||
database = DatabaseFactory.create({"content": "txtai.database.SQLite"})
|
||||
self.assertIsNotNone(database)
|
||||
|
||||
def testCustomBackendNotFound(self):
|
||||
"""
|
||||
Test resolving an unresolvable backend
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
DatabaseFactory.create({"content": "notfound.database"})
|
||||
@@ -0,0 +1,32 @@
|
||||
"""
|
||||
Database tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.database import Database
|
||||
|
||||
|
||||
class TestDatabase(unittest.TestCase):
|
||||
"""
|
||||
Base database tests.
|
||||
"""
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
database = Database({})
|
||||
|
||||
self.assertRaises(NotImplementedError, database.load, None)
|
||||
self.assertRaises(NotImplementedError, database.insert, None)
|
||||
self.assertRaises(NotImplementedError, database.delete, None)
|
||||
self.assertRaises(NotImplementedError, database.reindex, None)
|
||||
self.assertRaises(NotImplementedError, database.save, None)
|
||||
self.assertRaises(NotImplementedError, database.close)
|
||||
self.assertRaises(NotImplementedError, database.ids, None)
|
||||
self.assertRaises(NotImplementedError, database.count)
|
||||
self.assertRaises(NotImplementedError, database.resolve, None, None)
|
||||
self.assertRaises(NotImplementedError, database.embed, None, None)
|
||||
self.assertRaises(NotImplementedError, database.query, None, None, None, None)
|
||||
@@ -0,0 +1,84 @@
|
||||
"""
|
||||
DuckDB module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
from .testrdbms import Common
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestDuckDB(Common.TestRDBMS):
|
||||
"""
|
||||
Embeddings with content stored in DuckDB.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Content backend
|
||||
cls.backend = "duckdb"
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": cls.backend})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
@unittest.skipIf(os.name == "nt", "testArchive skipped on Windows")
|
||||
def testArchive(self):
|
||||
"""
|
||||
Test embeddings index archiving
|
||||
"""
|
||||
|
||||
super().testArchive()
|
||||
|
||||
def testFunction(self):
|
||||
"""
|
||||
Test custom functions
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"content": self.backend,
|
||||
"functions": [{"name": "textlength", "function": "testdatabase.testduckdb.length"}],
|
||||
}
|
||||
)
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("select textlength(text) length from txtai where id = 0", 1)[0]
|
||||
|
||||
self.assertEqual(int(result["length"]), 39)
|
||||
|
||||
|
||||
def length(text):
|
||||
"""
|
||||
Custom SQL function.
|
||||
"""
|
||||
|
||||
return len(text)
|
||||
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
Test encoding/decoding database objects
|
||||
"""
|
||||
|
||||
import glob
|
||||
import os
|
||||
import unittest
|
||||
import tempfile
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestEncoder(unittest.TestCase):
|
||||
"""
|
||||
Encoder tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = []
|
||||
for path in glob.glob(Utils.PATH + "/*jpg"):
|
||||
cls.data.append((path, {"object": Image.open(path)}, None))
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings(
|
||||
{"method": "sentence-transformers", "path": "sentence-transformers/clip-ViT-B-32", "content": True, "objects": "image"}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def testDefault(self):
|
||||
"""
|
||||
Test an index with default encoder
|
||||
"""
|
||||
|
||||
try:
|
||||
# Set default encoder
|
||||
self.embeddings.config["objects"] = True
|
||||
|
||||
# Test all database providers
|
||||
for content in ["duckdb", "sqlite"]:
|
||||
self.embeddings.config["content"] = content
|
||||
|
||||
data = [(0, {"object": bytearray([1, 2, 3]), "text": "default test"}, None)]
|
||||
|
||||
# Create an index
|
||||
self.embeddings.index(data)
|
||||
|
||||
result = self.embeddings.search("select object from txtai limit 1")[0]
|
||||
|
||||
self.assertEqual(result["object"].getvalue(), bytearray([1, 2, 3]))
|
||||
finally:
|
||||
self.embeddings.config["objects"] = "image"
|
||||
self.embeddings.config["content"] = True
|
||||
|
||||
def testImages(self):
|
||||
"""
|
||||
Test an index with image encoder
|
||||
"""
|
||||
|
||||
# Create an index for the list of images
|
||||
self.embeddings.index(self.data)
|
||||
|
||||
result = self.embeddings.search("select id, object from txtai where similar('universe') limit 1")[0]
|
||||
|
||||
self.assertTrue(result["id"].endswith("stars.jpg"))
|
||||
self.assertTrue(isinstance(result["object"], Image.Image))
|
||||
|
||||
@patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
|
||||
def testPickle(self):
|
||||
"""
|
||||
Test an index with pickle encoder
|
||||
"""
|
||||
|
||||
try:
|
||||
# Set pickle encoder
|
||||
self.embeddings.config["objects"] = "pickle"
|
||||
data = [(0, {"object": [1, 2, 3, 4, 5], "text": "default test"}, None)]
|
||||
|
||||
# Create an index
|
||||
self.embeddings.index(data)
|
||||
|
||||
result = self.embeddings.search("select object from txtai limit 1")[0]
|
||||
|
||||
self.assertEqual(result["object"], [1, 2, 3, 4, 5])
|
||||
finally:
|
||||
self.embeddings.config["objects"] = "image"
|
||||
|
||||
def testReindex(self):
|
||||
"""
|
||||
Test reindex with objects
|
||||
"""
|
||||
|
||||
# Create an index for the list of images
|
||||
self.embeddings.index(self.data)
|
||||
|
||||
# Reindex images
|
||||
self.embeddings.reindex({"method": "sentence-transformers", "path": "sentence-transformers/clip-ViT-B-32"})
|
||||
|
||||
result = self.embeddings.search("select id, object from txtai where similar('universe') limit 1")[0]
|
||||
|
||||
self.assertTrue(result["id"].endswith("stars.jpg"))
|
||||
self.assertTrue(isinstance(result["object"], Image.Image))
|
||||
|
||||
def testReindexFunction(self):
|
||||
"""
|
||||
Test reindex with objects and a function
|
||||
"""
|
||||
|
||||
try:
|
||||
# Streaming function that loads images on the fly
|
||||
def prepare(documents):
|
||||
for uid, data, tags in documents:
|
||||
yield (uid, Image.open(data), tags)
|
||||
|
||||
# Create an index for the list of images
|
||||
self.embeddings.index(self.data)
|
||||
|
||||
# Set default encoder and use function to load images
|
||||
self.embeddings.config["objects"] = True
|
||||
|
||||
# Save and load index to force default encoder
|
||||
index = os.path.join(tempfile.gettempdir(), "objects")
|
||||
self.embeddings.save(index)
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Reindex images
|
||||
self.embeddings.reindex({"method": "sentence-transformers", "path": "sentence-transformers/clip-ViT-B-32"}, function=prepare)
|
||||
|
||||
result = self.embeddings.search("select id, object from txtai where similar('universe') limit 1")[0]
|
||||
|
||||
self.assertTrue(result["id"].endswith("stars.jpg"))
|
||||
self.assertTrue(isinstance(result["object"], BytesIO))
|
||||
finally:
|
||||
self.embeddings.config["objects"] = "image"
|
||||
@@ -0,0 +1,935 @@
|
||||
"""
|
||||
Common file database module tests
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from txtai.embeddings import Embeddings, IndexNotFoundError
|
||||
from txtai.database import Embedded, RDBMS, SQLError
|
||||
|
||||
|
||||
class Common:
|
||||
"""
|
||||
Wraps common file database tests to prevent unit test discovery for this class.
|
||||
"""
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestRDBMS(unittest.TestCase):
|
||||
"""
|
||||
Embeddings with content stored in a file database tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Content backend
|
||||
cls.backend = None
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": cls.backend})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def testArchive(self):
|
||||
"""
|
||||
Test embeddings index archiving
|
||||
"""
|
||||
|
||||
for extension in ["tar.bz2", "tar.gz", "tar.xz", "zip"]:
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.{extension}")
|
||||
|
||||
self.embeddings.save(index)
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test offsets still work after save/load
|
||||
self.embeddings.upsert([(0, "Looking out into the dreadful abyss", None)])
|
||||
self.assertEqual(self.embeddings.count(), len(self.data))
|
||||
|
||||
def testAutoId(self):
|
||||
"""
|
||||
Test auto id generation
|
||||
"""
|
||||
|
||||
# Default sequence id
|
||||
embeddings = Embeddings(path="sentence-transformers/nli-mpnet-base-v2", content=self.backend)
|
||||
embeddings.index(self.data)
|
||||
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# UUID
|
||||
embeddings.config["autoid"] = "uuid4"
|
||||
embeddings.index(self.data)
|
||||
|
||||
result = embeddings.search(self.data[4], 1)[0]
|
||||
self.assertEqual(len(result["id"]), 36)
|
||||
|
||||
def testCheckpoint(self):
|
||||
"""
|
||||
Test embeddings index checkpoints
|
||||
"""
|
||||
|
||||
# Checkpoint directory
|
||||
checkpoint = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.checkpoint")
|
||||
|
||||
# Save embeddings checkpoint
|
||||
self.embeddings.index(self.data, checkpoint=checkpoint)
|
||||
|
||||
# Reindex with checkpoint
|
||||
self.embeddings.index(self.data, checkpoint=checkpoint)
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testColumns(self):
|
||||
"""
|
||||
Test custom text/object columns
|
||||
"""
|
||||
|
||||
embeddings = Embeddings({"keyword": True, "content": self.backend, "columns": {"text": "value"}})
|
||||
data = [{"value": x} for x in self.data]
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(data)])
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testClose(self):
|
||||
"""
|
||||
Test embeddings close
|
||||
"""
|
||||
|
||||
embeddings = None
|
||||
|
||||
# Create index twice to test open/close and ensure resources are freed
|
||||
for _ in range(2):
|
||||
embeddings = Embeddings(
|
||||
{"path": "sentence-transformers/nli-mpnet-base-v2", "scoring": {"method": "bm25", "terms": True}, "content": self.backend}
|
||||
)
|
||||
|
||||
# Add record to index
|
||||
embeddings.index([(0, "Close test", None)])
|
||||
|
||||
# Save index
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.close")
|
||||
embeddings.save(index)
|
||||
|
||||
# Close index
|
||||
embeddings.close()
|
||||
|
||||
# Test embeddings is empty
|
||||
self.assertIsNone(embeddings.ann)
|
||||
self.assertIsNone(embeddings.database)
|
||||
|
||||
def testData(self):
|
||||
"""
|
||||
Test content storage and retrieval
|
||||
"""
|
||||
|
||||
data = self.data + [{"date": "2021-01-01", "text": "Baby panda", "flag": 1}]
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[-1]["text"])
|
||||
|
||||
def testDelete(self):
|
||||
"""
|
||||
Test delete
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Delete best match
|
||||
self.embeddings.delete([4])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(self.embeddings.count(), 5)
|
||||
self.assertEqual(result["text"], self.data[5])
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test empty index
|
||||
"""
|
||||
|
||||
# Test search against empty index
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": self.backend})
|
||||
self.assertEqual(embeddings.search("test"), [])
|
||||
|
||||
# Test index with no data
|
||||
embeddings.index([])
|
||||
self.assertIsNone(embeddings.ann)
|
||||
|
||||
# Test upsert with no data
|
||||
embeddings.index([(0, "this is a test", None)])
|
||||
embeddings.upsert([])
|
||||
self.assertIsNotNone(embeddings.ann)
|
||||
|
||||
def testEmptyString(self):
|
||||
"""
|
||||
Test empty string indexing
|
||||
"""
|
||||
|
||||
# Test empty string
|
||||
self.embeddings.index([(0, "", None)])
|
||||
self.assertTrue(self.embeddings.search("test"))
|
||||
|
||||
# Test empty string with dict
|
||||
self.embeddings.index([(0, {"text": ""}, None)])
|
||||
self.assertTrue(self.embeddings.search("test"))
|
||||
|
||||
def testExplain(self):
|
||||
"""
|
||||
Test query explain
|
||||
"""
|
||||
|
||||
# Test explain with similarity
|
||||
result = self.embeddings.explain("feel good story", self.data)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
self.assertEqual(len(result.get("tokens")), 8)
|
||||
|
||||
def testExplainBatch(self):
|
||||
"""
|
||||
Test query explain batch
|
||||
"""
|
||||
|
||||
# Test explain with query
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
result = self.embeddings.batchexplain(["feel good story"], limit=1)[0][0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
self.assertEqual(len(result.get("tokens")), 8)
|
||||
|
||||
def testExplainEmpty(self):
|
||||
"""
|
||||
Test query explain with no filtering criteria
|
||||
"""
|
||||
|
||||
self.assertEqual(self.embeddings.explain("select * from txtai limit 1")[0]["id"], "0")
|
||||
|
||||
def testExpressions(self):
|
||||
"""
|
||||
Test expressions
|
||||
"""
|
||||
|
||||
# Test indexed expressions
|
||||
embeddings = Embeddings(
|
||||
path="sentence-transformers/nli-mpnet-base-v2",
|
||||
content=self.backend,
|
||||
expressions=[{"name": "textlength", "expression": "length(text)", "index": True}],
|
||||
)
|
||||
embeddings.index(self.data)
|
||||
|
||||
result = embeddings.search("SELECT textlength FROM txtai WHERE id = 0", 1)[0]
|
||||
self.assertEqual(result["textlength"], len(self.data[0]))
|
||||
|
||||
def testGenerator(self):
|
||||
"""
|
||||
Test index with a generator
|
||||
"""
|
||||
|
||||
def documents():
|
||||
for uid, text in enumerate(self.data):
|
||||
yield (uid, text, None)
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index(documents())
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testHybrid(self):
|
||||
"""
|
||||
Test hybrid search
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Index data with sparse + dense vectors.
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "hybrid": True, "content": self.backend})
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.hybrid")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Index data with sparse + dense vectors and unnormalized scores.
|
||||
embeddings.config["scoring"]["normalize"] = False
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Index data with sparse + dense vectors and bb25 normalized scores
|
||||
embeddings.config["scoring"]["normalize"] = "bb25"
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("canada intact iceberg a", 1)[0]
|
||||
self.assertEqual(result["text"], data[1][1])
|
||||
|
||||
# Test upsert
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test index
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testIndexTokens(self):
|
||||
"""
|
||||
Test index with tokens
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text.split(), None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testInfo(self):
|
||||
"""
|
||||
Test info
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
output = io.StringIO()
|
||||
with contextlib.redirect_stdout(output):
|
||||
self.embeddings.info()
|
||||
|
||||
self.assertIn("txtai", output.getvalue())
|
||||
|
||||
def testInstructions(self):
|
||||
"""
|
||||
Test indexing with instruction prefixes.
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"content": self.backend,
|
||||
"instructions": {"query": "query: ", "data": "passage: "},
|
||||
}
|
||||
)
|
||||
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testInvalidData(self):
|
||||
"""
|
||||
Test invalid JSON data
|
||||
"""
|
||||
|
||||
# Test invalid JSON value
|
||||
with self.assertRaises(ValueError):
|
||||
self.embeddings.index([(0, {"text": "This is a test", "flag": float("NaN")}, None)])
|
||||
|
||||
def testKeyword(self):
|
||||
"""
|
||||
Test keyword only (sparse) search
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Index data with sparse keyword vectors
|
||||
embeddings = Embeddings({"keyword": True, "content": self.backend})
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Test count method
|
||||
self.assertEqual(embeddings.count(), len(data))
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.keyword")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
def testMultiData(self):
|
||||
"""
|
||||
Test indexing with multiple data types (text, documents)
|
||||
"""
|
||||
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": self.backend, "batch": len(self.data)})
|
||||
|
||||
# Create an index using mixed data (text and documents)
|
||||
data = []
|
||||
for uid, text in enumerate(self.data):
|
||||
data.append((uid, text, None))
|
||||
data.append((uid, {"content": text}, None))
|
||||
|
||||
embeddings.index(data)
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testMultiSave(self):
|
||||
"""
|
||||
Test multiple successive saves
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Save original index
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.insert")
|
||||
self.embeddings.save(index)
|
||||
|
||||
# Modify index
|
||||
self.embeddings.upsert([(0, "Looking out into the dreadful abyss", None)])
|
||||
|
||||
# Save to a different location
|
||||
indexupdate = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.update")
|
||||
self.embeddings.save(indexupdate)
|
||||
|
||||
# Save to same location
|
||||
self.embeddings.save(index)
|
||||
|
||||
# Test all indexes match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
self.embeddings.load(index)
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
self.embeddings.load(indexupdate)
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testNoIndex(self):
|
||||
"""
|
||||
Test an embeddings instance with no available indexes
|
||||
"""
|
||||
|
||||
# Disable top-level indexing
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"content": self.backend,
|
||||
"defaults": False,
|
||||
}
|
||||
)
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
with self.assertRaises(IndexNotFoundError):
|
||||
embeddings.search("select id, text, score from txtai where similar('feel good story')")
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
db = RDBMS({})
|
||||
|
||||
self.assertRaises(NotImplementedError, db.connect, None)
|
||||
self.assertRaises(NotImplementedError, db.getcursor)
|
||||
self.assertRaises(NotImplementedError, db.jsonprefix)
|
||||
self.assertRaises(NotImplementedError, db.jsoncolumn, None)
|
||||
self.assertRaises(NotImplementedError, db.rows)
|
||||
self.assertRaises(NotImplementedError, db.addfunctions)
|
||||
|
||||
db = Embedded({})
|
||||
self.assertRaises(NotImplementedError, db.copy, None)
|
||||
|
||||
def testObject(self):
|
||||
"""
|
||||
Test object field
|
||||
"""
|
||||
|
||||
# Encode object
|
||||
embeddings = Embeddings({"defaults": False, "content": self.backend, "objects": True})
|
||||
embeddings.index([{"object": "binary data".encode("utf-8")}])
|
||||
|
||||
# Decode and test extracted object
|
||||
obj = embeddings.search("select object from txtai where id = 0")[0]["object"]
|
||||
self.assertEqual(str(obj.getvalue(), "utf-8"), "binary data")
|
||||
|
||||
@patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
|
||||
def testPickle(self):
|
||||
"""
|
||||
Test pickle configuration
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"format": "pickle",
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"content": self.backend,
|
||||
}
|
||||
)
|
||||
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.pickle")
|
||||
|
||||
embeddings.save(index)
|
||||
|
||||
# Check that config exists
|
||||
self.assertTrue(os.path.exists(os.path.join(index, "config")))
|
||||
|
||||
# Check that index can be reloaded
|
||||
embeddings.load(index)
|
||||
self.assertEqual(embeddings.count(), 6)
|
||||
|
||||
def testQuantize(self):
|
||||
"""
|
||||
Test scalar quantization
|
||||
"""
|
||||
|
||||
# Index data with 1-bit scalar quantization
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "quantize": 1, "content": self.backend})
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testQueryModel(self):
|
||||
"""
|
||||
Test index
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{"path": "sentence-transformers/nli-mpnet-base-v2", "content": self.backend, "query": {"path": "neuml/t5-small-txtsql"}}
|
||||
)
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story with win in text", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testReindex(self):
|
||||
"""
|
||||
Test reindex
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Delete records to test indexids still match
|
||||
self.embeddings.delete(([0, 1]))
|
||||
|
||||
# Reindex
|
||||
self.embeddings.reindex({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testSave(self):
|
||||
"""
|
||||
Test save
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}")
|
||||
|
||||
self.embeddings.save(index)
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test offsets still work after save/load
|
||||
self.embeddings.upsert([(0, "Looking out into the dreadful abyss", None)])
|
||||
self.assertEqual(self.embeddings.count(), len(self.data))
|
||||
|
||||
def testSettings(self):
|
||||
"""
|
||||
Test custom SQLite settings
|
||||
"""
|
||||
|
||||
# Index with write-ahead logging enabled
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": self.backend, "sqlite": {"wal": True}})
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testSQL(self):
|
||||
"""
|
||||
Test running a SQL query
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, {"text": text, "length": len(text), "attribute": f"ID{uid}"}, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Test similar
|
||||
result = self.embeddings.search(
|
||||
"select text, score from txtai where similar('feel good story') group by text, score having count(*) > 0 order by score desc", 1
|
||||
)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test similar with limits
|
||||
result = self.embeddings.search("select * from txtai where similar('feel good story', 1) limit 1")[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test similar with offset
|
||||
result = self.embeddings.search("select * from txtai where similar('feel good story') offset 1")[0]
|
||||
self.assertEqual(result["text"], self.data[5])
|
||||
|
||||
# Test where
|
||||
result = self.embeddings.search("select * from txtai where text like '%iceberg%'", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[1])
|
||||
|
||||
# Test count
|
||||
result = self.embeddings.search("select count(*) from txtai")[0]
|
||||
self.assertEqual(list(result.values())[0], len(self.data))
|
||||
|
||||
# Test columns
|
||||
result = self.embeddings.search("select id, text, length, data, entry from txtai")[0]
|
||||
self.assertEqual(sorted(result.keys()), ["data", "entry", "id", "length", "text"])
|
||||
|
||||
# Test column filtering
|
||||
result = self.embeddings.search("select text from txtai where attribute = 'ID4'", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test SQL parse error
|
||||
with self.assertRaises(SQLError):
|
||||
self.embeddings.search("select * from txtai where bad,query")
|
||||
|
||||
def testSQLBind(self):
|
||||
"""
|
||||
Test SQL statements with bind parameters
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Test similar clause bind parameters
|
||||
result = self.embeddings.search("select id, text, score from txtai where similar(:x)", parameters={"x": "feel good story"})[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test similar clause bind and non-bind parameters
|
||||
result = self.embeddings.search("select id, text, score from txtai where similar(:x, 0.5)", parameters={"x": "feel good story"})[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
# Test where filtering with bind parameters
|
||||
result = self.embeddings.search("select * from txtai where text like :x", parameters={"x": "%iceberg%"})[0]
|
||||
self.assertEqual(result["text"], self.data[1])
|
||||
|
||||
def testSparse(self):
|
||||
"""
|
||||
Test sparse vector search
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Index data with sparse vectors
|
||||
embeddings = Embeddings({"sparse": "sparse-encoder-testing/splade-bert-tiny-nq", "content": self.backend})
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Test count method
|
||||
self.assertEqual(embeddings.count(), len(data))
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.sparse")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
def testSubindex(self):
|
||||
"""
|
||||
Test subindex
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Disable top-level indexing and create subindex
|
||||
embeddings = Embeddings(
|
||||
{"content": self.backend, "defaults": False, "indexes": {"index1": {"path": "sentence-transformers/nli-mpnet-base-v2"}}}
|
||||
)
|
||||
embeddings.index(data)
|
||||
|
||||
# Test transform
|
||||
self.assertEqual(embeddings.transform("feel good story").shape, (768,))
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Run SQL search
|
||||
result = embeddings.search("select id, text, score from txtai where similar('feel good story', 10, 0.5)")[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Test missing index
|
||||
with self.assertRaises(IndexNotFoundError):
|
||||
embeddings.search("select id, text, score from txtai where similar('feel good story', 'notindex')")
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.subindex")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1])
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
# Check missing text is set to id when top-level indexing is disabled
|
||||
embeddings.upsert([(embeddings.count(), {"content": "empty text"}, None)])
|
||||
result = embeddings.search(f"{embeddings.count() - 1}", 1)[0]
|
||||
self.assertEqual(result["text"], str(embeddings.count() - 1))
|
||||
|
||||
# Close embeddings
|
||||
embeddings.close()
|
||||
|
||||
def testSubindexEmpty(self):
|
||||
"""
|
||||
Test loading an empty subindex
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, {"column1": text}, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Disable top-level indexing and create subindexes
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"content": self.backend,
|
||||
"defaults": False,
|
||||
"indexes": {
|
||||
"index1": {"path": "sentence-transformers/nli-mpnet-base-v2", "columns": {"text": "column1"}},
|
||||
"index2": {"path": "sentence-transformers/nli-mpnet-base-v2", "columns": {"text": "column2"}},
|
||||
},
|
||||
}
|
||||
)
|
||||
embeddings.index(data)
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), f"embeddings.{self.category()}.subindexempty")
|
||||
|
||||
# Save index
|
||||
embeddings.save(index)
|
||||
|
||||
# Test exists
|
||||
self.assertTrue(embeddings.exists(index))
|
||||
|
||||
# Load index
|
||||
embeddings.load(index)
|
||||
|
||||
# Test search
|
||||
result = embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[4][1]["text"])
|
||||
|
||||
def testTerms(self):
|
||||
"""
|
||||
Test extracting keyword terms from query
|
||||
"""
|
||||
|
||||
result = self.embeddings.terms("select * from txtai where similar('keyword terms')")
|
||||
self.assertEqual(result, "keyword terms")
|
||||
|
||||
def testTruncate(self):
|
||||
"""
|
||||
Test dimensionality truncation
|
||||
"""
|
||||
|
||||
# Truncate vectors to a specified number of dimensions
|
||||
embeddings = Embeddings(
|
||||
{"path": "sentence-transformers/nli-mpnet-base-v2", "dimensionality": 750, "content": self.backend, "vectors": {"revision": "main"}}
|
||||
)
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], self.data[4])
|
||||
|
||||
def testUpsert(self):
|
||||
"""
|
||||
Test upsert
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Reset embeddings for test
|
||||
self.embeddings.ann = None
|
||||
self.embeddings.database = None
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.upsert(data)
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
self.embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
|
||||
def testUpsertBatch(self):
|
||||
"""
|
||||
Test upsert batch
|
||||
"""
|
||||
|
||||
try:
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Reset embeddings for test
|
||||
self.embeddings.ann = None
|
||||
self.embeddings.database = None
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.upsert(data)
|
||||
|
||||
# Set batch size to 1
|
||||
self.embeddings.config["batch"] = 1
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
data[1] = (0, "Not good news", None)
|
||||
self.embeddings.upsert([data[0], data[1]])
|
||||
|
||||
# Search for best match
|
||||
result = self.embeddings.search("feel good story", 1)[0]
|
||||
|
||||
self.assertEqual(result["text"], data[0][1])
|
||||
finally:
|
||||
del self.embeddings.config["batch"]
|
||||
|
||||
def category(self):
|
||||
"""
|
||||
Content backend category.
|
||||
|
||||
Returns:
|
||||
category
|
||||
"""
|
||||
|
||||
return self.__class__.__name__.lower().replace("test", "")
|
||||
@@ -0,0 +1,325 @@
|
||||
"""
|
||||
SQL module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.database import DatabaseFactory, SQL, SQLError
|
||||
|
||||
|
||||
class TestSQL(unittest.TestCase):
|
||||
"""
|
||||
Test SQL parsing and generation.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
# Create SQL parser for SQLite
|
||||
cls.db = DatabaseFactory.create({"content": True})
|
||||
cls.db.initialize()
|
||||
|
||||
cls.sql = SQL(cls.db)
|
||||
|
||||
def testAlias(self):
|
||||
"""
|
||||
Test alias clauses
|
||||
"""
|
||||
|
||||
self.assertSql("select", "select a as a1 from txtai", "json_extract(data, '$.a') as a1")
|
||||
self.assertSql("select", "select a 'a1' from txtai", "json_extract(data, '$.a') 'a1'")
|
||||
self.assertSql("select", 'select a "a1" from txtai', "json_extract(data, '$.a') \"a1\"")
|
||||
self.assertSql("select", "select a a1 from txtai", "json_extract(data, '$.a') a1")
|
||||
self.assertSql(
|
||||
"select",
|
||||
"select a, b as b1, c, d + 1 as 'd1' from txtai",
|
||||
"json_extract(data, '$.a') as \"a\", json_extract(data, '$.b') as b1, "
|
||||
+ "json_extract(data, '$.c') as \"c\", json_extract(data, '$.d') + 1 as 'd1'",
|
||||
)
|
||||
self.assertSql("select", "select id as myid from txtai", "s.id as myid")
|
||||
self.assertSql("select", "select length(a) t from txtai", "length(json_extract(data, '$.a')) t")
|
||||
|
||||
self.assertSql("where", "select id as myid from txtai where myid != 3 and a != 1", "myid != 3 and json_extract(data, '$.a') != 1")
|
||||
self.assertSql("where", "select txt T from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt 'T' from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt \"T\" from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt as T from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt as 'T' from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
self.assertSql("where", "select txt as \"T\" from txtai where t LIKE '%abc%'", "t LIKE '%abc%'")
|
||||
|
||||
self.assertSql("groupby", "select id as myid, count(*) from txtai group by myid, a", "myid, json_extract(data, '$.a')")
|
||||
self.assertSql("orderby", "select id as myid from txtai order by myid, a", "myid, json_extract(data, '$.a')")
|
||||
|
||||
def testBadSQL(self):
|
||||
"""
|
||||
Test invalid SQL
|
||||
"""
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where order by")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where groupby order by")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where a(1)")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select a b c from txtai where id match id")
|
||||
|
||||
def testBracket(self):
|
||||
"""
|
||||
Test bracket expressions
|
||||
"""
|
||||
|
||||
self.assertSql("select", "select [a] from txtai", "json_extract(data, '$.a') as \"a\"")
|
||||
self.assertSql("select", "select [a] ab from txtai", "json_extract(data, '$.a') ab")
|
||||
self.assertSql("select", "select [abc] from txtai", "json_extract(data, '$.abc') as \"abc\"")
|
||||
self.assertSql("select", "select [id], text, score from txtai", "s.id, text, score")
|
||||
self.assertSql("select", "select [ab cd], text, score from txtai", "json_extract(data, '$.ab cd') as \"ab cd\", text, score")
|
||||
self.assertSql("select", "select [a[0]] from txtai", "json_extract(data, '$.a[0]') as \"a[0]\"")
|
||||
self.assertSql("select", "select [a[0].ab] from txtai", "json_extract(data, '$.a[0].ab') as \"a[0].ab\"")
|
||||
self.assertSql("select", "select [a[0].c[0]] from txtai", "json_extract(data, '$.a[0].c[0]') as \"a[0].c[0]\"")
|
||||
self.assertSql("select", "select avg([a]) from txtai", "avg(json_extract(data, '$.a')) as \"avg([a])\"")
|
||||
|
||||
# Test single quote escaping in bracket expressions
|
||||
self.assertSql("select", "select [field'] from txtai", "json_extract(data, '$.field''') as \"field'\"")
|
||||
|
||||
self.assertSql("where", "select * from txtai where [a b] < 1 or a > 1", "json_extract(data, '$.a b') < 1 or json_extract(data, '$.a') > 1")
|
||||
self.assertSql("where", "select [a[0].c[0]] a from txtai where a < 1", "a < 1")
|
||||
self.assertSql("groupby", "select * from txtai group by [a]", "json_extract(data, '$.a')")
|
||||
self.assertSql("orderby", "select * from txtai where order by [a]", "json_extract(data, '$.a')")
|
||||
|
||||
def testDistinct(self):
|
||||
"""
|
||||
Test distinct expressions
|
||||
"""
|
||||
|
||||
# Attributes
|
||||
self.assertSql("select", "select distinct id from txtai", "distinct s.id")
|
||||
self.assertSql("select", "select distinct id as myid from txtai", "distinct s.id as myid")
|
||||
self.assertSql("select", "select distinct a from txtai", "distinct json_extract(data, '$.a') as \"a\"")
|
||||
self.assertSql("select", "select distinct a.b from txtai", "distinct json_extract(data, '$.a.b') as \"a.b\"")
|
||||
|
||||
# Bracket expression
|
||||
self.assertSql("select", "select distinct [ab cd] from txtai", "distinct json_extract(data, '$.ab cd') as \"distinct[ab cd]\"")
|
||||
|
||||
# Function expression
|
||||
self.assertSql("select", "select distinct(id) from txtai", 'distinct(s.id) as "distinct(id)"')
|
||||
self.assertSql("select", "select count(distinct id) from txtai", 'count(distinct s.id) as "count(distinct id)"')
|
||||
self.assertSql("select", "select count(distinct a) from txtai", "count(distinct json_extract(data, '$.a')) as \"count(distinct a)\"")
|
||||
self.assertSql("select", "select count(distinct avg(id)) from txtai", 'count(distinct avg(s.id)) as "count(distinct avg(id))"')
|
||||
self.assertSql(
|
||||
"select", "select count(distinct avg(a)) from txtai", "count(distinct avg(json_extract(data, '$.a'))) as \"count(distinct avg(a))\""
|
||||
)
|
||||
|
||||
# Compound expression
|
||||
self.assertSql("select", "select distinct a/1 from txtai", "distinct json_extract(data, '$.a') / 1 as \"a / 1\"")
|
||||
self.assertSql("select", "select distinct(a/1) from txtai", "distinct(json_extract(data, '$.a') / 1) as \"distinct(a / 1)\"")
|
||||
|
||||
def testGroupby(self):
|
||||
"""
|
||||
Test group by clauses
|
||||
"""
|
||||
|
||||
prefix = "select count(*), flag from txtai "
|
||||
|
||||
self.assertSql("groupby", prefix + "group by text", "text")
|
||||
self.assertSql("groupby", prefix + "group by distinct(a)", "distinct(json_extract(data, '$.a'))")
|
||||
self.assertSql("groupby", prefix + "where a > 1 group by text", "text")
|
||||
|
||||
def testHaving(self):
|
||||
"""
|
||||
Test having clauses
|
||||
"""
|
||||
|
||||
prefix = "select count(*), flag from txtai "
|
||||
|
||||
self.assertSql("having", prefix + "group by text having count(*) > 1", "count(*) > 1")
|
||||
self.assertSql("having", prefix + "where flag = 1 group by text having count(*) > 1", "count(*) > 1")
|
||||
|
||||
def testIsSQL(self):
|
||||
"""
|
||||
Test SQL detection method.
|
||||
"""
|
||||
|
||||
self.assertTrue(self.sql.issql("select text from txtai where id = 1"))
|
||||
self.assertFalse(self.sql.issql(1234))
|
||||
|
||||
def testLimit(self):
|
||||
"""
|
||||
Test limit clauses
|
||||
"""
|
||||
|
||||
prefix = "select count(*) from txtai "
|
||||
|
||||
self.assertSql("limit", prefix + "limit 100", "100")
|
||||
|
||||
def testOffset(self):
|
||||
"""
|
||||
Test offset clauses
|
||||
"""
|
||||
|
||||
prefix = "select count(*) from txtai "
|
||||
|
||||
self.assertSql("offset", prefix + "limit 100 offset 50", "50")
|
||||
self.assertSql("offset", prefix + "offset 50", "50")
|
||||
|
||||
def testOrderby(self):
|
||||
"""
|
||||
Test order by clauses
|
||||
"""
|
||||
|
||||
prefix = "select * from txtai "
|
||||
|
||||
self.assertSql("orderby", prefix + "order by id", "s.id")
|
||||
self.assertSql("orderby", prefix + "order by id, text", "s.id, text")
|
||||
self.assertSql("orderby", prefix + "order by id asc", "s.id asc")
|
||||
self.assertSql("orderby", prefix + "order by id desc", "s.id desc")
|
||||
self.assertSql("orderby", prefix + "order by id asc, text desc", "s.id asc, text desc")
|
||||
|
||||
def testSelectBasic(self):
|
||||
"""
|
||||
Test basic select clauses
|
||||
"""
|
||||
|
||||
self.assertSql("select", "select id, indexid, tags from txtai", "s.id, s.indexid, s.tags")
|
||||
self.assertSql("select", "select id, indexid, flag from txtai", "s.id, s.indexid, json_extract(data, '$.flag') as \"flag\"")
|
||||
self.assertSql("select", "select id, indexid, a.b.c from txtai", "s.id, s.indexid, json_extract(data, '$.a.b.c') as \"a.b.c\"")
|
||||
self.assertSql("select", "select 'id', [id], (id) from txtai", "'id', s.id, (s.id)")
|
||||
self.assertSql("select", "select * from txtai", "*")
|
||||
|
||||
def testSelectCompound(self):
|
||||
"""
|
||||
Test compound select clauses
|
||||
"""
|
||||
|
||||
self.assertSql("select", "select a + 1 from txtai", "json_extract(data, '$.a') + 1 as \"a + 1\"")
|
||||
self.assertSql("select", "select 1 * a from txtai", "1 * json_extract(data, '$.a') as \"1 * a\"")
|
||||
self.assertSql("select", "select a/1 from txtai", "json_extract(data, '$.a') / 1 as \"a / 1\"")
|
||||
self.assertSql("select", "select avg(a-b) from txtai", "avg(json_extract(data, '$.a') - json_extract(data, '$.b')) as \"avg(a - b)\"")
|
||||
self.assertSql("select", "select distinct(text) from txtai", "distinct(text)")
|
||||
self.assertSql("select", "select id, score, (a/2)*3 from txtai", "s.id, score, (json_extract(data, '$.a') / 2) * 3 as \"(a / 2) * 3\"")
|
||||
self.assertSql("select", "select id, score, (a/2*3) from txtai", "s.id, score, (json_extract(data, '$.a') / 2 * 3) as \"(a / 2 * 3)\"")
|
||||
self.assertSql(
|
||||
"select",
|
||||
"select func(func2(indexid + 1), a) from txtai",
|
||||
"func(func2(s.indexid + 1), json_extract(data, '$.a')) as \"func(func2(indexid + 1), a)\"",
|
||||
)
|
||||
self.assertSql("select", "select func(func2(indexid + 1), a) a from txtai", "func(func2(s.indexid + 1), json_extract(data, '$.a')) a")
|
||||
self.assertSql("select", "select 'prefix' || id from txtai", "'prefix' || s.id as \"'prefix' || id\"")
|
||||
self.assertSql("select", "select 'prefix' || id id from txtai", "'prefix' || s.id id")
|
||||
self.assertSql("select", "select 'prefix' || a a from txtai", "'prefix' || json_extract(data, '$.a') a")
|
||||
|
||||
def testSimilar(self):
|
||||
"""
|
||||
Test similar functions
|
||||
"""
|
||||
|
||||
prefix = "select * from txtai "
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc')", "__SIMILAR__0")
|
||||
self.assertSql("similar", prefix + "where similar('abc')", [["abc"]])
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc') AND id = 1", "__SIMILAR__0 AND s.id = 1")
|
||||
self.assertSql("similar", prefix + "where similar('abc')", [["abc"]])
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc') and similar('def')", "__SIMILAR__0 and __SIMILAR__1")
|
||||
self.assertSql("similar", prefix + "where similar('abc') and similar('def')", [["abc"], ["def"]])
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc', 1000)", "__SIMILAR__0")
|
||||
self.assertSql("similar", prefix + "where similar('abc', 1000)", [["abc", "1000"]])
|
||||
|
||||
self.assertSql("where", prefix + "where similar('abc', 1000) and similar('def', 10)", "__SIMILAR__0 and __SIMILAR__1")
|
||||
self.assertSql("similar", prefix + "where similar('abc', 1000) and similar('def', 10)", [["abc", "1000"], ["def", "10"]])
|
||||
|
||||
self.assertSql("where", prefix + "where coalesce(similar('abc'), similar('abc'))", "coalesce(__SIMILAR__0, __SIMILAR__1)")
|
||||
self.assertSql("similar", prefix + "where coalesce(similar('abc'), similar('abc'))", [["abc"], ["abc"]])
|
||||
|
||||
def testUnterminated(self):
|
||||
"""
|
||||
Test unterminated clauses
|
||||
"""
|
||||
|
||||
# Unterminated bracket expressions
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select [a from txtai")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select avg([a) from txtai")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select [a[0] from txtai")
|
||||
|
||||
# Unterminated function expressions
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select func(a from txtai")
|
||||
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where coalesce(a")
|
||||
|
||||
# Unterminated similar clause
|
||||
with self.assertRaises(SQLError):
|
||||
self.db.search("select * from txtai where similar('abc'")
|
||||
|
||||
def testUpper(self):
|
||||
"""
|
||||
Test SQL statements are case insensitive.
|
||||
"""
|
||||
|
||||
self.assertSql("groupby", "SELECT * FROM TXTAI WHERE a = 1 GROUP BY id", "s.id")
|
||||
self.assertSql("orderby", "SELECT * FROM TXTAI WHERE a = 1 ORDER BY id", "s.id")
|
||||
|
||||
def testWhereBasic(self):
|
||||
"""
|
||||
Test basic where clauses
|
||||
"""
|
||||
|
||||
prefix = "select * from txtai "
|
||||
|
||||
self.assertSql("where", prefix + "where a = b", "json_extract(data, '$.a') = json_extract(data, '$.b')")
|
||||
self.assertSql("where", prefix + "where abc = def", "json_extract(data, '$.abc') = json_extract(data, '$.def')")
|
||||
self.assertSql("where", prefix + "where a = b.value", "json_extract(data, '$.a') = json_extract(data, '$.b.value')")
|
||||
self.assertSql("where", prefix + "where a = 1", "json_extract(data, '$.a') = 1")
|
||||
self.assertSql("where", prefix + "WHERE 1 = a", "1 = json_extract(data, '$.a')")
|
||||
self.assertSql("where", prefix + "WHERE a LIKE 'abc'", "json_extract(data, '$.a') LIKE 'abc'")
|
||||
self.assertSql("where", prefix + "WHERE a NOT LIKE 'abc'", "json_extract(data, '$.a') NOT LIKE 'abc'")
|
||||
self.assertSql("where", prefix + "WHERE a IN (1, 2, 3, b)", "json_extract(data, '$.a') IN (1, 2, 3, json_extract(data, '$.b'))")
|
||||
self.assertSql("where", prefix + "WHERE a is not null", "json_extract(data, '$.a') is not null")
|
||||
self.assertSql("where", prefix + "WHERE score >= 0.15", "score >= 0.15")
|
||||
|
||||
def testWhereCompound(self):
|
||||
"""
|
||||
Test compound where clauses
|
||||
"""
|
||||
|
||||
prefix = "select * from txtai "
|
||||
|
||||
self.assertSql("where", prefix + "where a > (b + 1)", "json_extract(data, '$.a') > (json_extract(data, '$.b') + 1)")
|
||||
self.assertSql("where", prefix + "where a > func('abc')", "json_extract(data, '$.a') > func('abc')")
|
||||
self.assertSql(
|
||||
"where", prefix + "where (id = 1 or id = 2) and a like 'abc'", "(s.id = 1 or s.id = 2) and json_extract(data, '$.a') like 'abc'"
|
||||
)
|
||||
self.assertSql(
|
||||
"where",
|
||||
prefix + "where a > f(d(b, c, 1),1)",
|
||||
"json_extract(data, '$.a') > f(d(json_extract(data, '$.b'), json_extract(data, '$.c'), 1), 1)",
|
||||
)
|
||||
self.assertSql("where", prefix + "where (id = 1 AND id = 2) OR indexid = 3", "(s.id = 1 AND s.id = 2) OR s.indexid = 3")
|
||||
self.assertSql("where", prefix + "where f(id) = b(id)", "f(s.id) = b(s.id)")
|
||||
self.assertSql("where", prefix + "WHERE f(id)", "f(s.id)")
|
||||
|
||||
def assertSql(self, clause, query, expected):
|
||||
"""
|
||||
Helper method to assert a query clause is as expected.
|
||||
|
||||
Args:
|
||||
clause: clause to select
|
||||
query: input query
|
||||
expected: expected transformed query value
|
||||
"""
|
||||
|
||||
self.assertEqual(self.sql(query)[clause], expected)
|
||||
@@ -0,0 +1,73 @@
|
||||
"""
|
||||
SQLite module tests
|
||||
"""
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
|
||||
from .testrdbms import Common
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestSQLite(Common.TestRDBMS):
|
||||
"""
|
||||
Embeddings with content stored in SQLite tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Content backend
|
||||
cls.backend = "sqlite"
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": cls.backend})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def testFunction(self):
|
||||
"""
|
||||
Test custom functions
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"content": self.backend,
|
||||
"functions": [{"name": "textlength", "function": "testdatabase.testsqlite.length"}],
|
||||
}
|
||||
)
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
result = embeddings.search("select textlength(text) length from txtai where id = 0", 1)[0]
|
||||
|
||||
self.assertEqual(int(result["length"]), 39)
|
||||
|
||||
|
||||
def length(text):
|
||||
"""
|
||||
Custom SQL function.
|
||||
"""
|
||||
|
||||
return len(text)
|
||||
@@ -0,0 +1,679 @@
|
||||
"""
|
||||
Embeddings module tests
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.embeddings import Embeddings, Reducer
|
||||
from txtai.serialize import SerializeFactory
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestEmbeddings(unittest.TestCase):
|
||||
"""
|
||||
Embeddings tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def testAutoId(self):
|
||||
"""
|
||||
Test auto id generation
|
||||
"""
|
||||
|
||||
# Default sequence id
|
||||
embeddings = Embeddings()
|
||||
embeddings.index(self.data)
|
||||
|
||||
uid = embeddings.search(self.data[4], 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# UUID
|
||||
embeddings = Embeddings(autoid="uuid4")
|
||||
embeddings.index(self.data)
|
||||
|
||||
uid = embeddings.search(self.data[4], 1)[0][0]
|
||||
self.assertEqual(len(uid), 36)
|
||||
|
||||
def testColumns(self):
|
||||
"""
|
||||
Test custom text/object columns
|
||||
"""
|
||||
|
||||
embeddings = Embeddings({"keyword": True, "columns": {"text": "value"}})
|
||||
data = [{"value": x} for x in self.data]
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(data)])
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("lottery", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testContext(self):
|
||||
"""
|
||||
Test embeddings context manager
|
||||
"""
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.context")
|
||||
|
||||
with Embeddings() as embeddings:
|
||||
embeddings.index(self.data)
|
||||
embeddings.save(index)
|
||||
|
||||
with Embeddings().load(index) as embeddings:
|
||||
uid = embeddings.search(self.data[4], 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testDefaults(self):
|
||||
"""
|
||||
Test default configuration
|
||||
"""
|
||||
|
||||
# Run index with no config which will fall back to default configuration
|
||||
embeddings = Embeddings()
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
self.assertEqual(embeddings.count(), 6)
|
||||
|
||||
def testDelete(self):
|
||||
"""
|
||||
Test delete
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Delete best match
|
||||
self.embeddings.delete([4])
|
||||
|
||||
# Search for best match
|
||||
uid = self.embeddings.search("feel good story", 1)[0][0]
|
||||
|
||||
self.assertEqual(self.embeddings.count(), 5)
|
||||
self.assertEqual(uid, 5)
|
||||
|
||||
def testDense(self):
|
||||
"""
|
||||
Test dense alias
|
||||
"""
|
||||
|
||||
# Dense flag is an alias for path
|
||||
embeddings = Embeddings(dense="sentence-transformers/nli-mpnet-base-v2")
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
self.assertEqual(embeddings.count(), 6)
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test empty index
|
||||
"""
|
||||
|
||||
# Test search against empty index
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
self.assertEqual(embeddings.search("test"), [])
|
||||
|
||||
# Test index with no data
|
||||
embeddings.index([])
|
||||
self.assertIsNone(embeddings.ann)
|
||||
|
||||
# Test upsert with no data
|
||||
embeddings.index([(0, "this is a test", None)])
|
||||
embeddings.upsert([])
|
||||
self.assertIsNotNone(embeddings.ann)
|
||||
|
||||
def testEmptyString(self):
|
||||
"""
|
||||
Test empty string indexing
|
||||
"""
|
||||
|
||||
# Test empty string
|
||||
self.embeddings.index([(0, "", None)])
|
||||
self.assertTrue(self.embeddings.search("test"))
|
||||
|
||||
# Test empty string with dict
|
||||
self.embeddings.index([(0, {"text": ""}, None)])
|
||||
self.assertTrue(self.embeddings.search("test"))
|
||||
|
||||
def testExternal(self):
|
||||
"""
|
||||
Test embeddings backed by external vectors
|
||||
"""
|
||||
|
||||
def transform(data):
|
||||
embeddings = []
|
||||
for text in data:
|
||||
# Create dummy embedding using sum and mean of character ordinals
|
||||
ordinals = [ord(c) for c in text]
|
||||
embeddings.append(np.array([sum(ordinals), np.mean(ordinals)]))
|
||||
|
||||
return embeddings
|
||||
|
||||
# Index data using simple embeddings transform method
|
||||
embeddings = Embeddings({"method": "external", "transform": transform})
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search(self.data[4], 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testExternalPrecomputed(self):
|
||||
"""
|
||||
Test embeddings backed by external pre-computed vectors
|
||||
"""
|
||||
|
||||
# Test with no transform function
|
||||
data = np.random.rand(5, 10).astype(np.float32)
|
||||
|
||||
embeddings = Embeddings({"method": "external"})
|
||||
embeddings.index([(uid, row, None) for uid, row in enumerate(data)])
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search(data[4], 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testHybrid(self):
|
||||
"""
|
||||
Test hybrid search
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Index data with sparse + dense vectors
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "hybrid": True})
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.hybrid")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Index data with sparse + dense vectors and unnormalized scores
|
||||
embeddings.config["scoring"]["normalize"] = False
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Index data with sparse + dense vectors and bb25 normalization
|
||||
embeddings.config["scoring"]["normalize"] = "bb25"
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("canada intact iceberg a", 1)[0][0]
|
||||
self.assertEqual(uid, 1)
|
||||
|
||||
# Test upsert
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 0)
|
||||
|
||||
def testIds(self):
|
||||
"""
|
||||
Test legacy config ids loading
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.ids")
|
||||
|
||||
# Save index
|
||||
self.embeddings.save(index)
|
||||
|
||||
# Set ids on config to simulate legacy ids format
|
||||
with open(f"{index}/config.json", "r", encoding="utf-8") as handle:
|
||||
config = json.load(handle)
|
||||
config["ids"] = list(range(len(self.data)))
|
||||
|
||||
with open(f"{index}/config.json", "w", encoding="utf-8") as handle:
|
||||
json.dump(config, handle, default=str, indent=2)
|
||||
|
||||
# Reload index
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
uid = self.embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Check that ids is not in config
|
||||
self.assertTrue("ids" not in self.embeddings.config)
|
||||
|
||||
@patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
|
||||
def testIdsPickle(self):
|
||||
"""
|
||||
Test legacy pickle ids
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.idspickle")
|
||||
|
||||
# Save index
|
||||
self.embeddings.save(index)
|
||||
|
||||
# Create ids as pickle
|
||||
path = os.path.join(tempfile.gettempdir(), "embeddings.idspickle", "ids")
|
||||
serializer = SerializeFactory.create("pickle", allowpickle=True)
|
||||
serializer.save(self.embeddings.ids.ids, path)
|
||||
|
||||
with self.assertWarns(RuntimeWarning):
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
uid = self.embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test index
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
uid = self.embeddings.search("feel good story", 1)[0][0]
|
||||
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testKeyword(self):
|
||||
"""
|
||||
Test keyword only (sparse) search
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Index data with sparse keyword vectors
|
||||
embeddings = Embeddings({"keyword": True})
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("lottery ticket", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Test count method
|
||||
self.assertEqual(embeddings.count(), len(data))
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.keyword")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("lottery ticket", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 0)
|
||||
|
||||
def testQuantize(self):
|
||||
"""
|
||||
Test scalar quantization
|
||||
"""
|
||||
|
||||
for ann in ["faiss", "numpy", "torch"]:
|
||||
# Index data with 1-bit scalar quantization
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "quantize": 1, "backend": ann})
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testReducer(self):
|
||||
"""
|
||||
Test reducer model
|
||||
"""
|
||||
|
||||
# Test model with single PCA component
|
||||
data = np.random.rand(5, 5).astype(np.float32)
|
||||
reducer = Reducer(data, 1)
|
||||
|
||||
# Generate query and keep original data to ensure it changes
|
||||
query = np.random.rand(1, 5).astype(np.float32)
|
||||
original = query.copy()
|
||||
|
||||
# Run test
|
||||
reducer(query)
|
||||
self.assertFalse(np.array_equal(query, original))
|
||||
|
||||
# Test model with multiple PCA components
|
||||
reducer = Reducer(data, 3)
|
||||
|
||||
# Generate query and keep original data to ensure it changes
|
||||
query = np.random.rand(5).astype(np.float32)
|
||||
original = query.copy()
|
||||
|
||||
# Run test
|
||||
reducer(query)
|
||||
self.assertFalse(np.array_equal(query, original))
|
||||
|
||||
@patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
|
||||
def testReducerLegacy(self):
|
||||
"""
|
||||
Test reducer model with legacy model format
|
||||
"""
|
||||
|
||||
# Test model with single PCA component
|
||||
data = np.random.rand(5, 5).astype(np.float32)
|
||||
reducer = Reducer(data, 1)
|
||||
|
||||
# Save legacy format
|
||||
path = os.path.join(tempfile.gettempdir(), "reducer")
|
||||
serializer = SerializeFactory.create("pickle", allowpickle=True)
|
||||
serializer.save(reducer.model, path)
|
||||
|
||||
# Load legacy format
|
||||
reducer = Reducer()
|
||||
reducer.load(path)
|
||||
|
||||
# Generate query and keep original data to ensure it changes
|
||||
query = np.random.rand(1, 5).astype(np.float32)
|
||||
original = query.copy()
|
||||
|
||||
# Run test
|
||||
reducer(query)
|
||||
self.assertFalse(np.array_equal(query, original))
|
||||
|
||||
def testSave(self):
|
||||
"""
|
||||
Test save
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.base")
|
||||
|
||||
self.embeddings.save(index)
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Search for best match
|
||||
uid = self.embeddings.search("feel good story", 1)[0][0]
|
||||
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Test offsets still work after save/load
|
||||
self.embeddings.upsert([(0, "Looking out into the dreadful abyss", None)])
|
||||
self.assertEqual(self.embeddings.count(), len(self.data))
|
||||
|
||||
def testShortcuts(self):
|
||||
"""
|
||||
Test embeddings creation shortcuts
|
||||
"""
|
||||
|
||||
tests = [
|
||||
({"keyword": True}, ["scoring"]),
|
||||
({"keyword": "sif"}, ["scoring"]),
|
||||
({"sparse": True}, ["scoring"]),
|
||||
({"dense": True}, ["ann"]),
|
||||
({"hybrid": True}, ["ann", "scoring"]),
|
||||
({"hybrid": "tfidf"}, ["ann", "scoring"]),
|
||||
({"hybrid": "sparse"}, ["ann", "scoring"]),
|
||||
({"graph": True}, ["graph"]),
|
||||
]
|
||||
|
||||
for config, checks in tests:
|
||||
embeddings = Embeddings(config)
|
||||
embeddings.index(["test"])
|
||||
|
||||
for attr in checks:
|
||||
self.assertIsNotNone(getattr(embeddings, attr))
|
||||
|
||||
def testSimilarity(self):
|
||||
"""
|
||||
Test similarity
|
||||
"""
|
||||
|
||||
# Get best matching id
|
||||
uid = self.embeddings.similarity("feel good story", self.data)[0][0]
|
||||
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testSparse(self):
|
||||
"""
|
||||
Test sparse vector search
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Index data with sparse vectors
|
||||
embeddings = Embeddings({"sparse": "sparse-encoder-testing/splade-bert-tiny-nq"})
|
||||
embeddings.index(data)
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("lottery ticket", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Test count method
|
||||
self.assertEqual(embeddings.count(), len(data))
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.sparse")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("lottery ticket", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Test similarity
|
||||
uid = embeddings.similarity("lottery ticket", self.data)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 0)
|
||||
|
||||
def testSubindex(self):
|
||||
"""
|
||||
Test subindex
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Disable top-level indexing and create subindex
|
||||
embeddings = Embeddings({"defaults": False, "indexes": {"index1": {"path": "sentence-transformers/nli-mpnet-base-v2"}}})
|
||||
embeddings.index(data)
|
||||
|
||||
# Test transform
|
||||
self.assertEqual(embeddings.transform("feel good story").shape, (768,))
|
||||
self.assertEqual(embeddings.transform("feel good story", index="index1").shape, (768,))
|
||||
with self.assertRaises(KeyError):
|
||||
embeddings.transform("feel good story", index="index2")
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.subindex")
|
||||
|
||||
# Test load/save
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Run search
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
uid = embeddings.search("feel good story", 10)[0][0]
|
||||
self.assertEqual(uid, 0)
|
||||
|
||||
# Check missing text is set to id when top-level indexing is disabled
|
||||
embeddings.upsert([(embeddings.count(), {"content": "empty text"}, None)])
|
||||
uid = embeddings.search(f"{embeddings.count() - 1}", 1)[0][0]
|
||||
self.assertEqual(uid, embeddings.count() - 1)
|
||||
|
||||
# Close embeddings
|
||||
embeddings.close()
|
||||
|
||||
def testTruncate(self):
|
||||
"""
|
||||
Test dimensionality truncation
|
||||
"""
|
||||
|
||||
# Truncate vectors to a specified number of dimensions
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "dimensionality": 750, "vectors": {"revision": "main"}})
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Search for best match
|
||||
uid = embeddings.search("feel good story", 1)[0][0]
|
||||
self.assertEqual(uid, 4)
|
||||
|
||||
def testUpsert(self):
|
||||
"""
|
||||
Test upsert
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
# Reset embeddings for test
|
||||
self.embeddings.ann = None
|
||||
self.embeddings.ids = None
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.upsert(data)
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: baby panda born", None)
|
||||
self.embeddings.upsert([data[0]])
|
||||
|
||||
# Search for best match
|
||||
uid = self.embeddings.search("feel good story", 1)[0][0]
|
||||
|
||||
self.assertEqual(uid, 0)
|
||||
|
||||
@patch("os.cpu_count")
|
||||
def testWords(self, cpucount):
|
||||
"""
|
||||
Test embeddings backed by word vectors
|
||||
"""
|
||||
|
||||
# Mock CPU count
|
||||
cpucount.return_value = 1
|
||||
|
||||
# Create dataset
|
||||
data = [(x, row.split(), None) for x, row in enumerate(self.data)]
|
||||
|
||||
# Create embeddings model, backed by word vectors
|
||||
embeddings = Embeddings({"path": "neuml/glove-6B-quantized", "scoring": "bm25", "pca": 3, "quantize": True})
|
||||
|
||||
# Call scoring and index methods
|
||||
embeddings.score(data)
|
||||
embeddings.index(data)
|
||||
|
||||
# Test search
|
||||
self.assertIsNotNone(embeddings.search("win", 1))
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "embeddings.wordvectors")
|
||||
|
||||
# Test save/load
|
||||
embeddings.save(index)
|
||||
embeddings.load(index)
|
||||
|
||||
# Test search
|
||||
self.assertIsNotNone(embeddings.search("win", 1))
|
||||
|
||||
@patch("os.cpu_count")
|
||||
def testWordsUpsert(self, cpucount):
|
||||
"""
|
||||
Test embeddings backed by word vectors with upserts
|
||||
"""
|
||||
|
||||
# Mock CPU count
|
||||
cpucount.return_value = 1
|
||||
|
||||
# Create dataset
|
||||
data = [(x, row.split(), None) for x, row in enumerate(self.data)]
|
||||
|
||||
# Create embeddings model, backed by word vectors
|
||||
embeddings = Embeddings({"path": "neuml/glove-6B/model.sqlite", "scoring": "bm25", "pca": 3})
|
||||
|
||||
# Call scoring and index methods
|
||||
embeddings.score(data)
|
||||
embeddings.index(data)
|
||||
|
||||
# Now upsert and override record
|
||||
data = [(0, "win win", None)]
|
||||
|
||||
# Update scoring and run upsert
|
||||
embeddings.score(data)
|
||||
embeddings.upsert(data)
|
||||
|
||||
# Test search after upsert
|
||||
uid = embeddings.search("win", 1)[0][0]
|
||||
self.assertEqual(uid, 0)
|
||||
@@ -0,0 +1,592 @@
|
||||
"""
|
||||
Graph module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import itertools
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from txtai.archive import ArchiveFactory
|
||||
from txtai.embeddings import Embeddings
|
||||
from txtai.graph import Graph, GraphFactory
|
||||
from txtai.serialize import SerializeFactory
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestGraph(unittest.TestCase):
|
||||
"""
|
||||
Graph tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
cls.config = {
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"content": True,
|
||||
"functions": [{"name": "graph", "function": "graph.attribute"}],
|
||||
"expressions": [
|
||||
{"name": "category", "expression": "graph(indexid, 'category')"},
|
||||
{"name": "topic", "expression": "graph(indexid, 'topic')"},
|
||||
{"name": "topicrank", "expression": "graph(indexid, 'topicrank')"},
|
||||
],
|
||||
"graph": {"limit": 5, "minscore": 0.2, "batchsize": 4, "approximate": False, "topics": {"categories": ["News"], "stopwords": ["the"]}},
|
||||
}
|
||||
|
||||
# Create embeddings instance
|
||||
cls.embeddings = Embeddings(cls.config)
|
||||
|
||||
def testAnalysis(self):
|
||||
"""
|
||||
Test analysis methods
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Graph centrality
|
||||
graph = self.embeddings.graph
|
||||
centrality = graph.centrality()
|
||||
self.assertEqual(list(centrality.keys())[0], 5)
|
||||
|
||||
# Page Rank
|
||||
pagerank = graph.pagerank()
|
||||
self.assertEqual(list(pagerank.keys())[0], 5)
|
||||
|
||||
# Path between nodes
|
||||
path = graph.showpath(4, 5)
|
||||
self.assertEqual(len(path), 2)
|
||||
|
||||
def testCommunity(self):
|
||||
"""
|
||||
Test community detection
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Get graph reference
|
||||
graph = self.embeddings.graph
|
||||
|
||||
# Rebuild topics with updated graph settings
|
||||
graph.config = {"topics": {"algorithm": "greedy"}}
|
||||
graph.addtopics()
|
||||
self.assertEqual(sum((len(graph.topics[x]) for x in graph.topics)), 6)
|
||||
|
||||
graph.config = {"topics": {"algorithm": "lpa"}}
|
||||
graph.addtopics()
|
||||
self.assertEqual(sum((len(graph.topics[x]) for x in graph.topics)), 4)
|
||||
|
||||
def testCustomBackend(self):
|
||||
"""
|
||||
Test resolving a custom backend
|
||||
"""
|
||||
|
||||
graph = GraphFactory.create({"backend": "txtai.graph.NetworkX"})
|
||||
graph.initialize()
|
||||
self.assertIsNotNone(graph)
|
||||
|
||||
def testCustomBackendNotFound(self):
|
||||
"""
|
||||
Test resolving an unresolvable backend
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
graph = GraphFactory.create({"backend": "notfound.graph"})
|
||||
graph.initialize()
|
||||
|
||||
def testDatabase(self):
|
||||
"""
|
||||
Test creating a Graph backed by a relational database
|
||||
"""
|
||||
|
||||
# Generate graph database
|
||||
path = os.path.join(tempfile.gettempdir(), "graph.sqlite")
|
||||
graph = GraphFactory.create({"backend": "rdbms", "url": f"sqlite:///{path}", "schema": "txtai"})
|
||||
|
||||
# Initialize the graph
|
||||
graph.initialize()
|
||||
|
||||
for x in range(5):
|
||||
graph.addnode(x, field=x)
|
||||
|
||||
for x, y in itertools.combinations(range(5), 2):
|
||||
graph.addedge(x, y)
|
||||
|
||||
# Test methods
|
||||
self.assertEqual(list(graph.scan()), [str(x) for x in range(5)])
|
||||
self.assertEqual(list(graph.scan(attribute="field")), [str(x) for x in range(5)])
|
||||
self.assertEqual(list(graph.filter([0]).scan()), [0])
|
||||
|
||||
# Test save/load
|
||||
graph.save(None)
|
||||
graph.load(None)
|
||||
self.assertEqual(list(graph.scan()), [str(x) for x in range(5)])
|
||||
|
||||
# Test remove node
|
||||
graph.delete([0])
|
||||
self.assertFalse(graph.hasnode(0))
|
||||
self.assertFalse(graph.hasedge(0))
|
||||
|
||||
# Close graph
|
||||
graph.close()
|
||||
|
||||
def testDefault(self):
|
||||
"""
|
||||
Test embeddings default graph setting
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(content=True, graph=True)
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
self.assertEqual(embeddings.graph.count(), len(self.data))
|
||||
|
||||
def testDelete(self):
|
||||
"""
|
||||
Test delete
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Delete row
|
||||
self.embeddings.delete([4])
|
||||
|
||||
# Validate counts
|
||||
graph = self.embeddings.graph
|
||||
self.assertEqual(graph.count(), 5)
|
||||
self.assertEqual(graph.edgecount(), 1)
|
||||
self.assertEqual(sum((len(graph.topics[x]) for x in graph.topics)), 5)
|
||||
self.assertEqual(len(graph.categories), 6)
|
||||
|
||||
def testEdges(self):
|
||||
"""
|
||||
Test edges
|
||||
"""
|
||||
|
||||
# Create graph
|
||||
graph = GraphFactory.create({})
|
||||
graph.initialize()
|
||||
graph.addedge(0, 1)
|
||||
|
||||
# Test edge exists
|
||||
self.assertTrue(graph.hasedge(0))
|
||||
self.assertTrue(graph.hasedge(0, 1))
|
||||
|
||||
def testFilter(self):
|
||||
"""
|
||||
Test creating filtered subgraphs
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Validate counts
|
||||
graph = self.embeddings.search("feel good story", graph=True)
|
||||
self.assertEqual(graph.count(), 3)
|
||||
self.assertEqual(graph.edgecount(), 2)
|
||||
|
||||
def testFunction(self):
|
||||
"""
|
||||
Test running graph functions with SQL
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Test function
|
||||
result = self.embeddings.search("select category, topic, topicrank from txtai where id = 0", 1)[0]
|
||||
|
||||
# Check columns have a value
|
||||
self.assertIsNotNone(result["category"])
|
||||
self.assertIsNotNone(result["topic"])
|
||||
self.assertIsNotNone(result["topicrank"])
|
||||
|
||||
def testFunctionReindex(self):
|
||||
"""
|
||||
Test running graph functions with SQL after reindex
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Test functions reset with a reindex
|
||||
self.embeddings.reindex(self.embeddings.config)
|
||||
|
||||
# Test function
|
||||
result = self.embeddings.search("select category, topic, topicrank from txtai where id = 0", 1)[0]
|
||||
|
||||
# Check columns have a value
|
||||
self.assertIsNotNone(result["category"])
|
||||
self.assertIsNotNone(result["topic"])
|
||||
self.assertIsNotNone(result["topicrank"])
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test index
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Validate counts
|
||||
graph = self.embeddings.graph
|
||||
self.assertEqual(graph.count(), 6)
|
||||
self.assertEqual(graph.edgecount(), 2)
|
||||
self.assertEqual(len(graph.topics), 6)
|
||||
self.assertEqual(len(graph.categories), 6)
|
||||
|
||||
@patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
|
||||
def testLegacy(self):
|
||||
"""
|
||||
Test loading a legacy graph in TAR format
|
||||
"""
|
||||
|
||||
# Create graph
|
||||
graph = GraphFactory.create({})
|
||||
graph.initialize()
|
||||
graph.addedge(0, 1)
|
||||
|
||||
categories = ["C1"]
|
||||
topics = {"T1": [0, 1]}
|
||||
|
||||
serializer = SerializeFactory.create("pickle", allowpickle=True)
|
||||
|
||||
# Save files to temporary directory and combine into TAR
|
||||
path = os.path.join(tempfile.gettempdir(), "graph.tar")
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
# Save graph
|
||||
serializer.save(graph.backend, f"{directory}/graph")
|
||||
|
||||
# Save categories, if necessary
|
||||
serializer.save(categories, f"{directory}/categories")
|
||||
|
||||
# Save topics, if necessary
|
||||
serializer.save(topics, f"{directory}/topics")
|
||||
|
||||
# Pack files
|
||||
archive = ArchiveFactory.create(directory)
|
||||
archive.save(path, "tar")
|
||||
|
||||
# Load loading legacy format
|
||||
graph = GraphFactory.create({})
|
||||
graph.load(path)
|
||||
|
||||
# Validate graph data is correct
|
||||
self.assertEqual(graph.count(), 2)
|
||||
self.assertEqual(graph.edgecount(), 1)
|
||||
self.assertEqual(graph.topics, topics)
|
||||
self.assertEqual(graph.categories, categories)
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
graph = Graph({})
|
||||
|
||||
self.assertRaises(NotImplementedError, graph.create)
|
||||
self.assertRaises(NotImplementedError, graph.count)
|
||||
self.assertRaises(NotImplementedError, graph.scan, None)
|
||||
self.assertRaises(NotImplementedError, graph.node, None)
|
||||
self.assertRaises(NotImplementedError, graph.addnode, None)
|
||||
self.assertRaises(NotImplementedError, graph.addnodes, None)
|
||||
self.assertRaises(NotImplementedError, graph.removenode, None)
|
||||
self.assertRaises(NotImplementedError, graph.hasnode, None)
|
||||
self.assertRaises(NotImplementedError, graph.attribute, None, None)
|
||||
self.assertRaises(NotImplementedError, graph.addattribute, None, None, None)
|
||||
self.assertRaises(NotImplementedError, graph.removeattribute, None, None)
|
||||
self.assertRaises(NotImplementedError, graph.edgecount)
|
||||
self.assertRaises(NotImplementedError, graph.edges, None)
|
||||
self.assertRaises(NotImplementedError, graph.addedge, None, None)
|
||||
self.assertRaises(NotImplementedError, graph.addedges, None)
|
||||
self.assertRaises(NotImplementedError, graph.hasedge, None, None)
|
||||
self.assertRaises(NotImplementedError, graph.centrality)
|
||||
self.assertRaises(NotImplementedError, graph.pagerank)
|
||||
self.assertRaises(NotImplementedError, graph.showpath, None, None)
|
||||
self.assertRaises(NotImplementedError, graph.isquery, None)
|
||||
self.assertRaises(NotImplementedError, graph.parse, None)
|
||||
self.assertRaises(NotImplementedError, graph.search, None)
|
||||
self.assertRaises(NotImplementedError, graph.communities, None)
|
||||
self.assertRaises(NotImplementedError, graph.load, None)
|
||||
self.assertRaises(NotImplementedError, graph.save, None)
|
||||
self.assertRaises(NotImplementedError, graph.loaddict, None)
|
||||
self.assertRaises(NotImplementedError, graph.savedict)
|
||||
|
||||
def testRelationships(self):
|
||||
"""
|
||||
Test manually-provided relationships
|
||||
"""
|
||||
|
||||
# Create relationships for id 0
|
||||
relationships = [{"id": f"ID{x}"} for x in range(1, len(self.data))]
|
||||
|
||||
# Test with content enabled
|
||||
self.embeddings.index({"id": f"ID{i}", "text": x, "relationships": relationships if i == 0 else None} for i, x in enumerate(self.data))
|
||||
self.assertEqual(len(self.embeddings.graph.edges(0)), len(self.data) - 1)
|
||||
|
||||
# Test with content disabled
|
||||
config = self.config.copy()
|
||||
config["content"] = False
|
||||
|
||||
embeddings = Embeddings(config)
|
||||
embeddings.index({"id": f"ID{i}", "text": x, "relationships": relationships if i == 0 else None} for i, x in enumerate(self.data))
|
||||
self.assertEqual(len(embeddings.graph.edges(0)), len(self.data) - 1)
|
||||
embeddings.close()
|
||||
|
||||
def testRelationshipsInvalid(self):
|
||||
"""
|
||||
Test manually-provided relationships with no matching id
|
||||
"""
|
||||
|
||||
# Create relationships for id 0
|
||||
relationships = [{"id": "INVALID"}]
|
||||
|
||||
# Index with invalid relationship
|
||||
self.embeddings.index({"text": x, "relationships": relationships if i == 0 else None} for i, x in enumerate(self.data))
|
||||
|
||||
# Validate only relationship is semantically-derived
|
||||
edges = list(self.embeddings.graph.edges(0))
|
||||
self.assertTrue(len(edges) == 1 and edges[0] != "INVALID")
|
||||
|
||||
def testResetTopics(self):
|
||||
"""
|
||||
Test resetting of topics
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(1, "text", None)])
|
||||
self.embeddings.upsert([(1, "graph", None)])
|
||||
self.assertEqual(list(self.embeddings.graph.topics.keys()), ["graph"])
|
||||
|
||||
def testSave(self):
|
||||
"""
|
||||
Test save
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "graph")
|
||||
|
||||
# Save and reload index
|
||||
self.embeddings.save(index)
|
||||
self.embeddings.load(index)
|
||||
|
||||
# Validate counts
|
||||
graph = self.embeddings.graph
|
||||
self.assertEqual(graph.count(), 6)
|
||||
self.assertEqual(graph.edgecount(), 2)
|
||||
self.assertEqual(sum((len(graph.topics[x]) for x in graph.topics)), 6)
|
||||
self.assertEqual(len(graph.categories), 6)
|
||||
|
||||
def testSaveDict(self):
|
||||
"""
|
||||
Test loading and saving to dictionaries
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Validate counts
|
||||
graph = self.embeddings.graph
|
||||
count, edgecount = graph.count(), graph.edgecount()
|
||||
|
||||
# Save and reload graph as dict
|
||||
data = graph.savedict()
|
||||
graph.loaddict(data)
|
||||
|
||||
# Validate counts
|
||||
self.assertEqual(graph.count(), count)
|
||||
self.assertEqual(graph.edgecount(), edgecount)
|
||||
|
||||
def testSearch(self):
|
||||
"""
|
||||
Test search
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Run standard search
|
||||
results = self.embeddings.search(
|
||||
"""
|
||||
MATCH (A)-[]->(B)
|
||||
RETURN A, B
|
||||
"""
|
||||
)
|
||||
self.assertEqual(len(results), 3)
|
||||
|
||||
# Run path search
|
||||
results = self.embeddings.search(
|
||||
"""
|
||||
MATCH P=()-[]->()
|
||||
RETURN P
|
||||
"""
|
||||
)
|
||||
self.assertEqual(len(results), 3)
|
||||
|
||||
# Run graph search
|
||||
g = self.embeddings.search(
|
||||
"""
|
||||
MATCH (A)-[]->(B)
|
||||
RETURN A, ID(B)
|
||||
""",
|
||||
graph=True,
|
||||
)
|
||||
self.assertEqual(g.count(), 3)
|
||||
|
||||
# Run path search
|
||||
results = self.embeddings.search(
|
||||
"""
|
||||
MATCH P=()-[]->()
|
||||
RETURN P
|
||||
""",
|
||||
graph=True,
|
||||
)
|
||||
self.assertEqual(g.count(), 3)
|
||||
|
||||
# Run similar search
|
||||
results = self.embeddings.search(
|
||||
"""
|
||||
MATCH P=(A)-[]->()
|
||||
WHERE SIMILAR(A, "feel good story")
|
||||
RETURN A
|
||||
ORDER BY A.score DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
graph=True,
|
||||
)
|
||||
self.assertEqual(list(results.scan())[0], 4)
|
||||
|
||||
def testSearchBatch(self):
|
||||
"""
|
||||
Test batch search
|
||||
"""
|
||||
|
||||
# Create an index for the list of text
|
||||
self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
# Run standard search
|
||||
results = self.embeddings.batchsearch(
|
||||
[
|
||||
"""
|
||||
MATCH (A)-[]->(B)
|
||||
RETURN A, B
|
||||
"""
|
||||
]
|
||||
)
|
||||
self.assertEqual(len(results[0]), 3)
|
||||
|
||||
def testSimple(self):
|
||||
"""
|
||||
Test creating a simple graph
|
||||
"""
|
||||
|
||||
graph = GraphFactory.create({"topics": {}})
|
||||
|
||||
# Initialize the graph
|
||||
graph.initialize()
|
||||
|
||||
for x in range(5):
|
||||
graph.addnode(x)
|
||||
|
||||
for x, y in itertools.combinations(range(5), 2):
|
||||
graph.addedge(x, y)
|
||||
|
||||
# Validate counts
|
||||
self.assertEqual(graph.count(), 5)
|
||||
self.assertEqual(graph.edgecount(), 10)
|
||||
|
||||
# Test missing edge
|
||||
self.assertIsNone(graph.edges(100))
|
||||
|
||||
# Test topics with no text
|
||||
graph.addtopics()
|
||||
self.assertEqual(len(graph.topics), 5)
|
||||
|
||||
def testSubindex(self):
|
||||
"""
|
||||
Test subindex
|
||||
"""
|
||||
|
||||
# Build data array
|
||||
data = [(uid, text, None) for uid, text in enumerate(self.data)]
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"content": True,
|
||||
"functions": [{"name": "graph", "function": "indexes.index1.graph.attribute"}],
|
||||
"expressions": [
|
||||
{"name": "category", "expression": "graph(indexid, 'category')"},
|
||||
{"name": "topic", "expression": "graph(indexid, 'topic')"},
|
||||
{"name": "topicrank", "expression": "graph(indexid, 'topicrank')"},
|
||||
],
|
||||
"indexes": {
|
||||
"index1": {
|
||||
"path": "sentence-transformers/nli-mpnet-base-v2",
|
||||
"graph": {
|
||||
"limit": 5,
|
||||
"minscore": 0.2,
|
||||
"batchsize": 4,
|
||||
"approximate": False,
|
||||
"topics": {"categories": ["News"], "stopwords": ["the"]},
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
# Create an index for the list of text
|
||||
embeddings.index(data)
|
||||
|
||||
# Test function
|
||||
result = embeddings.search("select id, category, topic, topicrank from txtai where id = 0", 1)[0]
|
||||
|
||||
# Check columns have a value
|
||||
self.assertIsNotNone(result["category"])
|
||||
self.assertIsNotNone(result["topic"])
|
||||
self.assertIsNotNone(result["topicrank"])
|
||||
|
||||
# Update data
|
||||
data[0] = (0, "Feel good story: lottery winner announced", None)
|
||||
embeddings.upsert([data[0]])
|
||||
|
||||
# Test function
|
||||
result = embeddings.search("select id, category, topic, topicrank from txtai where id = 0", 1)[0]
|
||||
|
||||
# Check columns have a value
|
||||
self.assertIsNotNone(result["category"])
|
||||
self.assertIsNotNone(result["topic"])
|
||||
self.assertIsNotNone(result["topicrank"])
|
||||
|
||||
def testUpsert(self):
|
||||
"""
|
||||
Test upsert
|
||||
"""
|
||||
|
||||
# Update data
|
||||
self.embeddings.upsert([(0, {"text": "Canadian ice shelf collapses".split()}, None)])
|
||||
|
||||
# Validate counts
|
||||
graph = self.embeddings.graph
|
||||
self.assertEqual(graph.count(), 6)
|
||||
self.assertEqual(graph.edgecount(), 2)
|
||||
self.assertEqual(sum((len(graph.topics[x]) for x in graph.topics)), 6)
|
||||
self.assertEqual(len(graph.categories), 6)
|
||||
@@ -0,0 +1,102 @@
|
||||
"""
|
||||
Library module tests
|
||||
"""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
# pylint: disable=C0415,W0611,W0621
|
||||
import txtai
|
||||
|
||||
|
||||
class TestLibrary(unittest.TestCase):
|
||||
"""
|
||||
Simulates core libraries not being installed.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Simulates core libraries not being installed
|
||||
"""
|
||||
|
||||
modules = [
|
||||
"huggingface_hub",
|
||||
"huggingface_hub.errors",
|
||||
"numpy",
|
||||
"regex",
|
||||
"safetensors",
|
||||
"transformers",
|
||||
"transformers.configuration_utils",
|
||||
"transformers.modeling_utils",
|
||||
"transformers.modeling_outputs",
|
||||
"torch",
|
||||
"torch.nn",
|
||||
"torch.onnx",
|
||||
"torch.utils.data",
|
||||
"yaml",
|
||||
]
|
||||
|
||||
# Get handle to all currently loaded txtai modules
|
||||
modules = modules + [key for key in sys.modules if key.startswith("txtai")]
|
||||
cls.modules = {module: None for module in modules}
|
||||
|
||||
# Replace loaded modules with stubs. Save modules for later reloading
|
||||
for module in cls.modules:
|
||||
if module in sys.modules:
|
||||
cls.modules[module] = sys.modules[module]
|
||||
|
||||
# Remove txtai modules. Set optional dependencies to None to prevent reloading.
|
||||
if "txtai" in module:
|
||||
if module in sys.modules:
|
||||
del sys.modules[module]
|
||||
else:
|
||||
sys.modules[module] = None
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Resets modules environment back to initial state.
|
||||
"""
|
||||
|
||||
# Reset replaced modules in setup
|
||||
for key, value in cls.modules.items():
|
||||
if value:
|
||||
sys.modules[key] = value
|
||||
else:
|
||||
del sys.modules[key]
|
||||
|
||||
def testLibrary(self):
|
||||
"""
|
||||
Test core libraries not being installed
|
||||
"""
|
||||
|
||||
# pylint: disable=W0106
|
||||
from txtai.util import Library
|
||||
|
||||
lib = Library()
|
||||
|
||||
# Test transformers stubs
|
||||
for x in [lib.arguments(), lib.config(), lib.dataset(), lib.hferror(), lib.module(), lib.model()]:
|
||||
self.assertTrue(x.__module__.endswith("library"))
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
lib.huggingface_hub().hf_hub_download
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
lib.numpy().dot
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
lib.regex().compile
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
lib.safetensors().safe_open
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
lib.torch().device
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
lib.transformers().AutoModel
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
lib.yaml().safe_open
|
||||
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
Models module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from txtai.models import Models
|
||||
|
||||
|
||||
class TestModels(unittest.TestCase):
|
||||
"""
|
||||
Models tests.
|
||||
"""
|
||||
|
||||
@patch("torch.cuda.is_available")
|
||||
def testDeviceid(self, cuda):
|
||||
"""
|
||||
Test the deviceid method
|
||||
"""
|
||||
|
||||
cuda.return_value = True
|
||||
self.assertEqual(Models.deviceid(True), 0)
|
||||
self.assertEqual(Models.deviceid(False), -1)
|
||||
self.assertEqual(Models.deviceid(0), 0)
|
||||
self.assertEqual(Models.deviceid(1), 1)
|
||||
|
||||
# Test direct torch device
|
||||
# pylint: disable=E1101
|
||||
self.assertEqual(Models.deviceid(torch.device("cpu")), torch.device("cpu"))
|
||||
|
||||
cuda.return_value = False
|
||||
self.assertEqual(Models.deviceid(True), -1)
|
||||
self.assertEqual(Models.deviceid(False), -1)
|
||||
self.assertEqual(Models.deviceid(0), -1)
|
||||
self.assertEqual(Models.deviceid(1), -1)
|
||||
|
||||
def testDevice(self):
|
||||
"""
|
||||
Test the device method
|
||||
"""
|
||||
|
||||
# pylint: disable=E1101
|
||||
self.assertEqual(Models.device("cpu"), torch.device("cpu"))
|
||||
self.assertEqual(Models.device(torch.device("cpu")), torch.device("cpu"))
|
||||
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
Pooling module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.models import Models, ClsPooling, LastPooling, MeanPooling, PoolingFactory
|
||||
|
||||
|
||||
class TestPooling(unittest.TestCase):
|
||||
"""
|
||||
Pooling tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize device
|
||||
"""
|
||||
|
||||
# Device id
|
||||
cls.device = Models.deviceid(True)
|
||||
|
||||
def testCLS(self):
|
||||
"""
|
||||
Test CLS pooling
|
||||
"""
|
||||
|
||||
# Test CLS pooling
|
||||
pooling = PoolingFactory.create({"path": "flax-sentence-embeddings/multi-qa_v1-MiniLM-L6-cls_dot", "device": self.device})
|
||||
self.assertEqual(type(pooling), ClsPooling)
|
||||
|
||||
pooling = PoolingFactory.create({"method": "clspooling", "path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device})
|
||||
self.assertEqual(type(pooling), ClsPooling)
|
||||
|
||||
# Test CLS pooling encoding
|
||||
self.assertEqual(pooling.encode(["test"])[0].shape, (768,))
|
||||
|
||||
def testLast(self):
|
||||
"""
|
||||
Test last pooling
|
||||
"""
|
||||
|
||||
# Test last pooling
|
||||
pooling = PoolingFactory.create({"path": "neuml/bert-tiny-sts-last-pooling", "device": self.device})
|
||||
self.assertEqual(type(pooling), LastPooling)
|
||||
|
||||
pooling = PoolingFactory.create({"method": "lastpooling", "path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device})
|
||||
self.assertEqual(type(pooling), LastPooling)
|
||||
|
||||
# Test last pooling encoding
|
||||
self.assertEqual(pooling.encode(["test"])[0].shape, (768,))
|
||||
|
||||
def testLength(self):
|
||||
"""
|
||||
Test pooling with max_seq_length
|
||||
"""
|
||||
|
||||
# Test reading max_seq_length parmaeter
|
||||
pooling = PoolingFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device, "maxlength": True})
|
||||
self.assertEqual(pooling.maxlength, 75)
|
||||
|
||||
# Test specified maxlength
|
||||
pooling = PoolingFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device, "maxlength": 256})
|
||||
self.assertEqual(pooling.maxlength, 256)
|
||||
|
||||
# Test max_seq_length is ignored when parameter is omitted
|
||||
pooling = PoolingFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device})
|
||||
self.assertEqual(pooling.maxlength, 512)
|
||||
|
||||
# Test maxlength when max_seq_length not present
|
||||
pooling = PoolingFactory.create({"path": "hf-internal-testing/tiny-random-gpt2", "device": self.device, "maxlength": True})
|
||||
self.assertEqual(pooling.maxlength, 1024)
|
||||
|
||||
def testMean(self):
|
||||
"""
|
||||
Test mean pooling
|
||||
"""
|
||||
|
||||
# Test mean pooling
|
||||
pooling = PoolingFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2", "device": self.device})
|
||||
self.assertEqual(type(pooling), MeanPooling)
|
||||
|
||||
pooling = PoolingFactory.create(
|
||||
{"method": "meanpooling", "path": "flax-sentence-embeddings/multi-qa_v1-MiniLM-L6-cls_dot", "device": self.device}
|
||||
)
|
||||
self.assertEqual(type(pooling), MeanPooling)
|
||||
|
||||
def testMuvera(self):
|
||||
"""
|
||||
Test late pooling with MUVERA fixed dimensional encoding
|
||||
"""
|
||||
|
||||
# Test MUVERA encoding
|
||||
for model in ["neuml/colbert-bert-tiny", "neuml/pylate-bert-tiny"]:
|
||||
# Test defaults
|
||||
pooling = PoolingFactory.create({"path": model, "device": self.device})
|
||||
self.assertEqual(pooling.encode(["test"], category="query").shape, (1, 10240))
|
||||
|
||||
# Test custom settings
|
||||
pooling = PoolingFactory.create(
|
||||
{"path": model, "device": self.device, "modelargs": {"muvera": {"repetitions": 5, "hashes": 2, "projection": 8}}}
|
||||
)
|
||||
self.assertEqual(pooling.encode(["test"], category="data").shape, (1, 160))
|
||||
|
||||
def testPrompts(self):
|
||||
"""
|
||||
Test instruction prompts
|
||||
"""
|
||||
|
||||
# Load model with prompts
|
||||
pooling = PoolingFactory.create({"path": "neuml/bert-tiny-prompts", "device": self.device, "loadprompts": True})
|
||||
|
||||
# Test prompts are prepended
|
||||
self.assertEqual(pooling.preencode(["abc"], "query")[0], "query: abc")
|
||||
self.assertEqual(pooling.preencode(["text"], "data")[0], "document: text")
|
||||
|
||||
# Load model with prompts disabled (default)
|
||||
pooling = PoolingFactory.create({"path": "neuml/bert-tiny-prompts", "device": self.device})
|
||||
|
||||
# Test that prompts are not prepended
|
||||
self.assertEqual(pooling.preencode(["abc"], "query")[0], "abc")
|
||||
self.assertEqual(pooling.preencode(["text"], "data")[0], "text")
|
||||
@@ -0,0 +1,414 @@
|
||||
"""
|
||||
Optional module tests
|
||||
"""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
# pylint: disable=C0415,W0611,W0621
|
||||
import timm
|
||||
import txtai
|
||||
|
||||
|
||||
class TestOptional(unittest.TestCase):
|
||||
"""
|
||||
Optional tests. Simulates optional dependencies not being installed.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Simulate optional packages not being installed
|
||||
"""
|
||||
|
||||
modules = [
|
||||
"ai_edge_litert.compiled_model",
|
||||
"annoy",
|
||||
"bitsandbytes",
|
||||
"bs4",
|
||||
"chonkie",
|
||||
"croniter",
|
||||
"docling.document_converter",
|
||||
"duckdb",
|
||||
"faiss",
|
||||
"fastapi",
|
||||
"ggml",
|
||||
"gliner",
|
||||
"grandcypher",
|
||||
"grand",
|
||||
"hnswlib",
|
||||
"httpx",
|
||||
"imagehash",
|
||||
"libcloud.storage.providers",
|
||||
"litellm",
|
||||
"liteparse",
|
||||
"litert_lm",
|
||||
"llama_cpp",
|
||||
"model2vec",
|
||||
"msgpack",
|
||||
"networkx",
|
||||
"nltk",
|
||||
"onnxmltools",
|
||||
"onnxruntime",
|
||||
"onnxruntime.quantization",
|
||||
"pandas",
|
||||
"peft",
|
||||
"pgvector",
|
||||
"PIL",
|
||||
"rich",
|
||||
"scipy",
|
||||
"scipy.sparse",
|
||||
"sentence_transformers",
|
||||
"sklearn.decomposition",
|
||||
"smolagents",
|
||||
"sounddevice",
|
||||
"soundfile",
|
||||
"sqlalchemy",
|
||||
"sqlite_vec",
|
||||
"staticvectors",
|
||||
"tika",
|
||||
"ttstokenizer",
|
||||
"turbovec",
|
||||
"xmltodict",
|
||||
]
|
||||
|
||||
# Get handle to all currently loaded txtai modules
|
||||
modules = modules + [key for key in sys.modules if key.startswith("txtai")]
|
||||
cls.modules = {module: None for module in modules}
|
||||
|
||||
# Replace loaded modules with stubs. Save modules for later reloading
|
||||
for module in cls.modules:
|
||||
if module in sys.modules:
|
||||
cls.modules[module] = sys.modules[module]
|
||||
|
||||
# Remove txtai modules. Set optional dependencies to None to prevent reloading.
|
||||
if "txtai" in module:
|
||||
if module in sys.modules:
|
||||
del sys.modules[module]
|
||||
else:
|
||||
sys.modules[module] = None
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Resets modules environment back to initial state.
|
||||
"""
|
||||
|
||||
# Reset replaced modules in setup
|
||||
for key, value in cls.modules.items():
|
||||
if value:
|
||||
sys.modules[key] = value
|
||||
else:
|
||||
del sys.modules[key]
|
||||
|
||||
def testAgent(self):
|
||||
"""
|
||||
Test missing agent dependencies
|
||||
"""
|
||||
|
||||
from txtai.agent import Agent
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Agent(llm="hf-internal-testing/tiny-random-LlamaForCausalLM", max_steps=1)
|
||||
|
||||
def testANN(self):
|
||||
"""
|
||||
Test missing ANN dependencies
|
||||
"""
|
||||
|
||||
from txtai.ann import ANNFactory, SparseANNFactory
|
||||
|
||||
# Test dense methods
|
||||
with self.assertRaises(ImportError):
|
||||
ANNFactory.create({"backend": "annoy"})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ANNFactory.create({"backend": "faiss"})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ANNFactory.create({"backend": "ggml"})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ANNFactory.create({"backend": "hnsw"})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ANNFactory.create({"backend": "pgvector"})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ANNFactory.create({"backend": "sqlite"})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ANNFactory.create({"backend": "torch", "torch": {"quantize": True}})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ANNFactory.create({"backend": "turbovec"})
|
||||
|
||||
# Test sparse methods
|
||||
with self.assertRaises(ImportError):
|
||||
SparseANNFactory.create({"backend": "ivfsparse"})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
SparseANNFactory.create({"backend": "pgsparse"})
|
||||
|
||||
def testApi(self):
|
||||
"""
|
||||
Test missing api dependencies
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
import txtai.api
|
||||
|
||||
def testConsole(self):
|
||||
"""
|
||||
Test missing console dependencies
|
||||
"""
|
||||
|
||||
from txtai.console import Console
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Console()
|
||||
|
||||
def testCloud(self):
|
||||
"""
|
||||
Test missing cloud dependencies
|
||||
"""
|
||||
|
||||
from txtai.cloud import ObjectStorage
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ObjectStorage(None)
|
||||
|
||||
def testDatabase(self):
|
||||
"""
|
||||
Test missing database dependencies
|
||||
"""
|
||||
|
||||
from txtai.database import Client, DuckDB, ImageEncoder
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Client({})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
DuckDB({})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ImageEncoder()
|
||||
|
||||
def testGraph(self):
|
||||
"""
|
||||
Test missing graph dependencies
|
||||
"""
|
||||
|
||||
from txtai.graph import GraphFactory, Query
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
GraphFactory.create({"backend": "networkx"})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
GraphFactory.create({"backend": "rdbms"})
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Query()
|
||||
|
||||
def testModel(self):
|
||||
"""
|
||||
Test missing model dependencies
|
||||
"""
|
||||
|
||||
from txtai.embeddings import Reducer
|
||||
from txtai.models import OnnxModel
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Reducer()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
OnnxModel(None)
|
||||
|
||||
# pylint: disable=R0915
|
||||
def testPipeline(self):
|
||||
"""
|
||||
Test missing pipeline dependencies
|
||||
"""
|
||||
|
||||
from txtai.pipeline import (
|
||||
AudioMixer,
|
||||
AudioStream,
|
||||
Caption,
|
||||
Entity,
|
||||
FileToHTML,
|
||||
HFOnnx,
|
||||
HFTrainer,
|
||||
HTMLToMarkdown,
|
||||
ImageHash,
|
||||
LiteLLM,
|
||||
LiteRT,
|
||||
LlamaCpp,
|
||||
Microphone,
|
||||
MLOnnx,
|
||||
Objects,
|
||||
OpenCode,
|
||||
Segmentation,
|
||||
Tabular,
|
||||
TextToAudio,
|
||||
TextToSpeech,
|
||||
Transcription,
|
||||
Translation,
|
||||
)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
AudioMixer()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
AudioStream()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Caption()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Entity("neuml/gliner-bert-tiny")
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
FileToHTML(backend="docling")
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
FileToHTML(backend="liteparse")
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
FileToHTML(backend="tika")
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
HFOnnx()("google/bert_uncased_L-2_H-128_A-2", quantize=True)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
HFTrainer()(None, None, lora=True)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
HTMLToMarkdown()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ImageHash()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
LiteLLM("huggingface/t5-small")
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
LiteRT("litert-community/SmolLM2-360M-Instruct/SmolLM2_360M_instruct.litertlm")
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
LlamaCpp("TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF/tinyllama-1.1b-chat-v0.3.Q2_K.gguf")
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Microphone()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
MLOnnx()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Objects()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
OpenCode("opencode")
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Segmentation(sentences=True)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Segmentation(chunker="token")
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Tabular()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
TextToAudio()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
TextToSpeech()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Transcription()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Translation().detect(["test"])
|
||||
|
||||
def testSerialize(self):
|
||||
"""
|
||||
Test missing msgpack dependency
|
||||
"""
|
||||
|
||||
from txtai.serialize import MessagePack
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
MessagePack()
|
||||
|
||||
def testScoring(self):
|
||||
"""
|
||||
Test missing scoring dependencies
|
||||
"""
|
||||
|
||||
from txtai.scoring import ScoringFactory
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ScoringFactory.create({"method": "pgtext"})
|
||||
|
||||
def testVectors(self):
|
||||
"""
|
||||
Test missing vector dependencies
|
||||
"""
|
||||
|
||||
from txtai.vectors import SparseVectors, VectorsFactory, SparseVectorsFactory
|
||||
from txtai.util import SparseArray
|
||||
|
||||
# Test dense vectors
|
||||
with self.assertRaises(ImportError):
|
||||
VectorsFactory.create({"method": "litellm", "path": "huggingface/sentence-transformers/all-MiniLM-L6-v2"}, None)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
VectorsFactory.create({"method": "litert", "path": "neuml/bert-hash-nano-embeddings-litert/bert-hash-nano-embeddings-int4.tflite"}, None)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
VectorsFactory.create({"method": "llama.cpp", "path": "nomic-ai/nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.Q2_K.gguf"}, None)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
VectorsFactory.create({"method": "model2vec", "path": "minishlab/M2V_base_output"}, None)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
VectorsFactory.create({"method": "sentence-transformers", "path": "sentence-transformers/nli-mpnet-base-v2"}, None)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
VectorsFactory.create({"method": "words"}, None)
|
||||
|
||||
# Test default model
|
||||
model = VectorsFactory.create({"path": "sentence-transformers/all-MiniLM-L6-v2"}, None)
|
||||
self.assertIsNotNone(model)
|
||||
|
||||
# Test sparse vectors
|
||||
with self.assertRaises(ImportError):
|
||||
SparseVectors(None, None, None)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
SparseVectorsFactory.create({"method": "sentence-transformers", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}, None)
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
SparseArray()
|
||||
|
||||
def testWorkflow(self):
|
||||
"""
|
||||
Test missing workflow dependencies
|
||||
"""
|
||||
|
||||
from txtai.workflow import ExportTask, ImageTask, ServiceTask, StorageTask, Workflow
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ExportTask()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ImageTask()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ServiceTask()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
StorageTask()
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
Workflow([], workers=1).schedule(None, [])
|
||||
@@ -0,0 +1,29 @@
|
||||
"""
|
||||
AudioMixer module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.pipeline import AudioMixer
|
||||
|
||||
|
||||
class TestAudioStream(unittest.TestCase):
|
||||
"""
|
||||
AudioStream tests.
|
||||
"""
|
||||
|
||||
def testAudioStream(self):
|
||||
"""
|
||||
Test mixing audio streams
|
||||
"""
|
||||
|
||||
audio1 = np.random.rand(2, 5000), 100
|
||||
audio2 = np.random.rand(2, 5000), 100
|
||||
|
||||
mixer = AudioMixer()
|
||||
audio, rate = mixer((audio1, audio2))
|
||||
|
||||
self.assertEqual(audio.shape, (2, 5000))
|
||||
self.assertEqual(rate, 100)
|
||||
@@ -0,0 +1,37 @@
|
||||
"""
|
||||
AudioStream module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import soundfile as sf
|
||||
|
||||
from txtai.pipeline import AudioStream
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestAudioStream(unittest.TestCase):
|
||||
"""
|
||||
AudioStream tests.
|
||||
"""
|
||||
|
||||
@patch("sounddevice.play")
|
||||
def testAudioStream(self, play):
|
||||
"""
|
||||
Test playing audio
|
||||
"""
|
||||
|
||||
play.return_value = True
|
||||
|
||||
# Read audio data
|
||||
audio, rate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
|
||||
|
||||
stream = AudioStream()
|
||||
self.assertIsNotNone(stream([(audio, rate), AudioStream.COMPLETE]))
|
||||
|
||||
# Wait for completion
|
||||
stream.wait()
|
||||
@@ -0,0 +1,81 @@
|
||||
"""
|
||||
Microphone module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
from txtai.pipeline import Microphone
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestMicrophone(unittest.TestCase):
|
||||
"""
|
||||
Microphone tests.
|
||||
"""
|
||||
|
||||
# pylint: disable=C0115,C0116
|
||||
@patch("sounddevice.RawInputStream")
|
||||
def testMicrophone(self, inputstream):
|
||||
"""
|
||||
Test listening to microphone
|
||||
"""
|
||||
|
||||
class RawInputStream:
|
||||
def __init__(self, **kwargs):
|
||||
self.args = kwargs
|
||||
|
||||
# Read audio data
|
||||
self.index, self.passes = 0, 0
|
||||
audio, self.samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
|
||||
|
||||
# Convert data to PCM
|
||||
self.audio = self.int16(audio)
|
||||
|
||||
# Start with random data to test that speech is not detected
|
||||
self.data = np.concatenate((self.audio * 50, np.zeros(shape=self.audio.shape, dtype=np.int16)))
|
||||
|
||||
def start(self):
|
||||
pass
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
def read(self, size):
|
||||
# Get chunk
|
||||
chunk = self.data[self.index : self.index + size]
|
||||
self.index += size
|
||||
|
||||
# Initial pass is random data, 2nd pass is speech data
|
||||
if self.index > len(self.data):
|
||||
if not self.passes:
|
||||
self.index, self.passes = 0, self.passes + 1
|
||||
self.data = self.audio
|
||||
elif self.index >= len(self.audio) * 10:
|
||||
# Break out of loop if speech continues to not be detected
|
||||
raise IOError("Data exhausted")
|
||||
|
||||
return chunk, False
|
||||
|
||||
def int16(self, data):
|
||||
i = np.iinfo(np.int16)
|
||||
absmax = 2 ** (i.bits - 1)
|
||||
offset = i.min + absmax
|
||||
return (data * absmax + offset).clip(i.min, i.max).astype(np.int16)
|
||||
|
||||
# Mock input stream
|
||||
inputstream.side_effect = RawInputStream
|
||||
|
||||
# Create microphone pipeline and read data
|
||||
pipeline = Microphone()
|
||||
data, rate = pipeline()
|
||||
|
||||
# Validate sample rate and length of data
|
||||
self.assertEqual(len(data), 91220)
|
||||
self.assertEqual(rate, 16000)
|
||||
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
TextToAudio module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import TextToAudio
|
||||
|
||||
|
||||
class TestTextToAudio(unittest.TestCase):
|
||||
"""
|
||||
TextToAudio tests.
|
||||
"""
|
||||
|
||||
def testTextToAudio(self):
|
||||
"""
|
||||
Test generating audio for text
|
||||
"""
|
||||
|
||||
tta = TextToAudio("hf-internal-testing/tiny-random-MusicgenForConditionalGeneration")
|
||||
|
||||
# Check that data is generated
|
||||
audio, rate = tta("This is a test")
|
||||
|
||||
self.assertGreater(len(audio), 0)
|
||||
self.assertEqual(rate, 24000)
|
||||
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
TextToSpeech module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from txtai.pipeline import TextToSpeech
|
||||
|
||||
|
||||
class TestTextToSpeech(unittest.TestCase):
|
||||
"""
|
||||
TextToSpeech tests.
|
||||
"""
|
||||
|
||||
def testESPnet(self):
|
||||
"""
|
||||
Test generating speech for text with an ESPnet model
|
||||
"""
|
||||
|
||||
tts = TextToSpeech()
|
||||
|
||||
# Check that data is generated
|
||||
speech, rate = tts("This is a test")
|
||||
|
||||
self.assertGreater(len(speech), 0)
|
||||
self.assertEqual(rate, 22050)
|
||||
|
||||
def testKokoro(self):
|
||||
"""
|
||||
Test generating speech for text with a Kokoro model
|
||||
"""
|
||||
|
||||
tts = TextToSpeech("neuml/kokoro-int8-onnx", maxtokens=2)
|
||||
|
||||
# Check that data is generated
|
||||
speech, rate = tts("This is a test")
|
||||
|
||||
self.assertGreater(len(speech), 0)
|
||||
self.assertEqual(rate, 22050)
|
||||
|
||||
@patch("onnxruntime.get_available_providers")
|
||||
@patch("torch.cuda.is_available")
|
||||
def testProviders(self, cuda, providers):
|
||||
"""
|
||||
Test that GPU provider is detected
|
||||
"""
|
||||
|
||||
# Test CUDA and onnxruntime-gpu installed
|
||||
cuda.return_value = True
|
||||
providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
|
||||
tts = TextToSpeech()
|
||||
self.assertEqual(tts.providers()[0][0], "CUDAExecutionProvider")
|
||||
|
||||
def testSpeechT5(self):
|
||||
"""
|
||||
Test generating speech for text with a SpeechT5 model
|
||||
"""
|
||||
|
||||
tts = TextToSpeech("neuml/txtai-speecht5-onnx")
|
||||
|
||||
# Check that data is generated
|
||||
speech, rate = tts("This is a test")
|
||||
|
||||
self.assertGreater(len(speech), 0)
|
||||
self.assertEqual(rate, 22050)
|
||||
|
||||
def testStreaming(self):
|
||||
"""
|
||||
Test streaming speech generation
|
||||
"""
|
||||
|
||||
tts = TextToSpeech()
|
||||
|
||||
# Check that data is generated
|
||||
speech, rate = list(tts("This is a test. And another".split(), stream=True))[0]
|
||||
|
||||
# Check that data is generated
|
||||
self.assertGreater(len(speech), 0)
|
||||
self.assertEqual(rate, 22050)
|
||||
@@ -0,0 +1,104 @@
|
||||
"""
|
||||
Transcription module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
from scipy import signal
|
||||
|
||||
from txtai.pipeline import Transcription
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestTranscription(unittest.TestCase):
|
||||
"""
|
||||
Transcription tests.
|
||||
"""
|
||||
|
||||
def testArray(self):
|
||||
"""
|
||||
Test audio data to text transcription
|
||||
"""
|
||||
|
||||
transcribe = Transcription()
|
||||
|
||||
# Read audio data
|
||||
raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
|
||||
|
||||
self.assertEqual(transcribe((raw, samplerate)), "Make huge profits without working make up to one hundred thousand dollars a day")
|
||||
self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")
|
||||
|
||||
def testChunks(self):
|
||||
"""
|
||||
Test splitting transcription into chunks
|
||||
"""
|
||||
|
||||
transcribe = Transcription()
|
||||
|
||||
result = transcribe(Utils.PATH + "/Make_huge_profits.wav", join=False)[0]
|
||||
|
||||
self.assertIsInstance(result["raw"], np.ndarray)
|
||||
self.assertIsNotNone(result["rate"])
|
||||
self.assertEqual(result["text"], "Make huge profits without working make up to one hundred thousand dollars a day")
|
||||
|
||||
def testFile(self):
|
||||
"""
|
||||
Test audio file to text transcription
|
||||
"""
|
||||
|
||||
transcribe = Transcription()
|
||||
|
||||
self.assertEqual(
|
||||
transcribe(Utils.PATH + "/Make_huge_profits.wav"), "Make huge profits without working make up to one hundred thousand dollars a day"
|
||||
)
|
||||
|
||||
def testGenerateArguments(self):
|
||||
"""
|
||||
Test transcription with generation keyword arguments
|
||||
"""
|
||||
|
||||
transcribe = Transcription()
|
||||
|
||||
# Read audio data
|
||||
raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
|
||||
|
||||
self.assertEqual(
|
||||
transcribe(raw, samplerate, language="English", task="transcribe"),
|
||||
"Make huge profits without working make up to one hundred thousand dollars a day",
|
||||
)
|
||||
|
||||
def testResample(self):
|
||||
"""
|
||||
Test resampled audio file to text transcription
|
||||
"""
|
||||
|
||||
transcribe = Transcription()
|
||||
|
||||
# Read audio data
|
||||
raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
|
||||
|
||||
# Resample for testing
|
||||
samples = round(len(raw) * float(22050) / samplerate)
|
||||
raw, samplerate = signal.resample(raw, samples), 22050
|
||||
|
||||
self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")
|
||||
|
||||
def testStereo(self):
|
||||
"""
|
||||
Test audio file in stereo to text transcription
|
||||
"""
|
||||
|
||||
transcribe = Transcription()
|
||||
|
||||
# Read audio data
|
||||
raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
|
||||
|
||||
# Convert mono to stereo
|
||||
raw = np.column_stack((raw, raw))
|
||||
|
||||
self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")
|
||||
@@ -0,0 +1,24 @@
|
||||
"""
|
||||
FileToHTML module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from txtai.pipeline.data.filetohtml import Tika
|
||||
|
||||
|
||||
class TestFileToHTML(unittest.TestCase):
|
||||
"""
|
||||
FileToHTML tests.
|
||||
"""
|
||||
|
||||
@patch.dict(os.environ, {"TIKA_JAVA": "1112444abc"})
|
||||
def testTika(self):
|
||||
"""
|
||||
Test the Tika.available returns False when Java is not available
|
||||
"""
|
||||
|
||||
self.assertFalse(Tika.available())
|
||||
+123
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
Tabular module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Tabular
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestTabular(unittest.TestCase):
|
||||
"""
|
||||
Tabular tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create single tabular instance
|
||||
"""
|
||||
|
||||
cls.tabular = Tabular("id", ["text"])
|
||||
|
||||
def testContent(self):
|
||||
"""
|
||||
Test parsing additional content
|
||||
"""
|
||||
|
||||
tabular = Tabular("id", ["text"], True)
|
||||
|
||||
row = {"id": 0, "text": "This is a test", "flag": 1}
|
||||
|
||||
# When content is enabled, both (uid, text, tags) and (uid, data, tags) rows are generated
|
||||
# given that data doesn't necessarily include the text to index
|
||||
rows = tabular([row])
|
||||
uid, data, _ = rows[1]
|
||||
|
||||
# Data should contain the entire input row
|
||||
self.assertEqual(uid, 0)
|
||||
self.assertEqual(data, row)
|
||||
|
||||
# Only select flag field
|
||||
tabular.content = ["flag"]
|
||||
rows = tabular([row])
|
||||
uid, data, _ = rows[1]
|
||||
|
||||
# Data should only contain a single field, flag
|
||||
self.assertEqual(uid, 0)
|
||||
self.assertTrue(list(data.keys()) == ["flag"])
|
||||
self.assertEqual(data["flag"], 1)
|
||||
|
||||
def testCSV(self):
|
||||
"""
|
||||
Test parsing a CSV file
|
||||
"""
|
||||
|
||||
rows = self.tabular([Utils.PATH + "/tabular.csv"])
|
||||
uid, text, _ = rows[0][0]
|
||||
|
||||
self.assertEqual(uid, 0)
|
||||
self.assertEqual(text, "The first sentence")
|
||||
|
||||
def testDict(self):
|
||||
"""
|
||||
Test parsing a dict
|
||||
"""
|
||||
|
||||
rows = self.tabular([{"id": 0, "text": "This is a test"}])
|
||||
uid, text, _ = rows[0]
|
||||
|
||||
self.assertEqual(uid, 0)
|
||||
self.assertEqual(text, "This is a test")
|
||||
|
||||
def testInvalid(self):
|
||||
"""
|
||||
Test invalid file paths
|
||||
"""
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
self.tabular([Utils.PATH + "/article.pdf"])
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
self.tabular(["https://invalid.path"])
|
||||
|
||||
def testList(self):
|
||||
"""
|
||||
Test parsing a list
|
||||
"""
|
||||
|
||||
rows = self.tabular([[{"id": 0, "text": "This is a test"}]])
|
||||
uid, text, _ = rows[0][0]
|
||||
|
||||
self.assertEqual(uid, 0)
|
||||
self.assertEqual(text, "This is a test")
|
||||
|
||||
def testMissingColumns(self):
|
||||
"""
|
||||
Test rows with uneven or missing columns
|
||||
"""
|
||||
|
||||
tabular = Tabular("id", ["text"], True)
|
||||
|
||||
rows = tabular([{"id": 0, "text": "This is a test", "metadata": "meta"}, {"id": 1, "text": "This is a test"}])
|
||||
|
||||
# When content is enabled both (id, text, tag) and (id, data, tag) tuples are generated given that
|
||||
# data doesn't necessarily include the text to index
|
||||
_, data, _ = rows[3]
|
||||
|
||||
self.assertIsNone(data["metadata"])
|
||||
|
||||
def testNoColumns(self):
|
||||
"""
|
||||
Test creating text without specifying columns
|
||||
"""
|
||||
|
||||
tabular = Tabular("id")
|
||||
rows = tabular([{"id": 0, "text": "This is a test", "summary": "Describes text in more detail"}])
|
||||
uid, text, _ = rows[0]
|
||||
|
||||
self.assertEqual(uid, 0)
|
||||
self.assertEqual(text, "This is a test. Describes text in more detail")
|
||||
+253
@@ -0,0 +1,253 @@
|
||||
"""
|
||||
Textractor module tests
|
||||
"""
|
||||
|
||||
import platform
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Textractor
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestTextractor(unittest.TestCase):
|
||||
"""
|
||||
Textractor tests.
|
||||
"""
|
||||
|
||||
def testClean(self):
|
||||
"""
|
||||
Test text cleaning method
|
||||
"""
|
||||
|
||||
# Default text cleaning
|
||||
textractor = Textractor()
|
||||
self.assertEqual(textractor(" a b c "), "a b c")
|
||||
|
||||
# Require text to be minlength
|
||||
textractor = Textractor(minlength=10)
|
||||
self.assertEqual(textractor(" a b c "), None)
|
||||
|
||||
# Disable text cleaning
|
||||
textractor = Textractor(cleantext=False, minlength=10)
|
||||
self.assertEqual(textractor(" a b c "), " a b c ")
|
||||
|
||||
def testChonkie(self):
|
||||
"""
|
||||
Test a chonkie chunker
|
||||
"""
|
||||
|
||||
# Test chonkie chunking
|
||||
textractor = Textractor(chunker="sentence", chunk_size=5, chunk_overlap=0)
|
||||
self.assertEqual(textractor("This is a test. And another test."), ["This is a test.", "And another test."])
|
||||
|
||||
# Test bad chunker throws an exception
|
||||
with self.assertRaises(AttributeError):
|
||||
textractor = Textractor(chunker="badchunker")
|
||||
|
||||
def testDefault(self):
|
||||
"""
|
||||
Test default text extraction
|
||||
"""
|
||||
|
||||
# Text input
|
||||
textractor = Textractor(backend=None)
|
||||
text = textractor(Utils.PATH + "/tabular.csv")
|
||||
self.assertEqual(len(text), 125)
|
||||
|
||||
# Markdown input
|
||||
textractor = Textractor(sections=True)
|
||||
sections = textractor("# Heading 1\nText1\n\n# Heading 2\nText2\n")
|
||||
|
||||
# Check number of sections is as expected
|
||||
self.assertEqual(len(sections), 2)
|
||||
|
||||
@unittest.skipIf(platform.system() == "Darwin", "Docling skipped on macOS to avoid MPS issues")
|
||||
def testDocling(self):
|
||||
"""
|
||||
Test docling backend
|
||||
"""
|
||||
|
||||
textractor = Textractor(backend="docling")
|
||||
|
||||
# Extract text and check for Markdown formatting
|
||||
text = textractor(Utils.PATH + "/article.pdf")
|
||||
self.assertTrue("## Introducing txtai" in text)
|
||||
|
||||
def testLines(self):
|
||||
"""
|
||||
Test extraction to lines
|
||||
"""
|
||||
|
||||
textractor = Textractor(lines=True)
|
||||
|
||||
# Extract text as lines
|
||||
lines = textractor(Utils.PATH + "/article.pdf")
|
||||
|
||||
# Check number of lines is as expected
|
||||
self.assertEqual(len(lines), 35)
|
||||
|
||||
def testLiteParse(self):
|
||||
"""
|
||||
Test liteparse backend
|
||||
"""
|
||||
|
||||
textractor = Textractor(backend="liteparse")
|
||||
|
||||
# Extract text and check for Markdown formatting
|
||||
text = textractor(Utils.PATH + "/article.pdf")
|
||||
self.assertTrue("# Introducing txtai" in text)
|
||||
|
||||
def testHTML(self):
|
||||
"""
|
||||
Test HTML to Markdown
|
||||
"""
|
||||
|
||||
# Headings
|
||||
self.assertMarkdown("<h1>This is a test</h1>", "# This is a test")
|
||||
self.assertMarkdown("<h6>This is a test</h6>", "###### This is a test")
|
||||
|
||||
# Blockquotes
|
||||
self.assertMarkdown("<blockquote>This is a test</blockquote>", "> This is a test")
|
||||
|
||||
# Lists
|
||||
self.assertMarkdown("<ul><li>Test1</li><li>Test2</li></ul>", "- Test1\n- Test2")
|
||||
self.assertMarkdown("<ol><li>Test1</li><li>Test2</li></ol>", "1. Test1\n2. Test2")
|
||||
|
||||
# Code
|
||||
self.assertMarkdown("<code>This is a test</code>", "```\nThis is a test\n```")
|
||||
self.assertMarkdown("<pre>This is a test</pre>", "```\nThis is a test\n```")
|
||||
|
||||
# Tables
|
||||
self.assertMarkdown(
|
||||
"<table><tr><th>Header1</th><th>Header2</th></tr><tr><td>Test1</td><td>Test2</td></tr></table>",
|
||||
"|Header1|Header2|\n|---|---|\n|Test1|Test2|",
|
||||
)
|
||||
|
||||
# Ignore list
|
||||
self.assertMarkdown("<aside>This is a test</aside>", "")
|
||||
|
||||
# Text formatting
|
||||
self.assertMarkdown("<p>This is a test</p>", "This is a test")
|
||||
self.assertMarkdown("<p>This is a <b>test</b</p>", "This is a **test**")
|
||||
self.assertMarkdown("<p>This is a <strong>test</strong></p>", "This is a **test**")
|
||||
self.assertMarkdown("<p>This is a <i>test</i></p>", "This is a *test*")
|
||||
self.assertMarkdown("<p>This is a <em>test</em></p>", "This is a *test*")
|
||||
self.assertMarkdown("<p>This is a <a href='link'>test</a>", "This is a [test](link)")
|
||||
|
||||
# Collapse to outer tag
|
||||
self.assertMarkdown("<p>This is a <strong><em>test</em></strong></p>", "This is a **test**")
|
||||
self.assertMarkdown("<p>This is a <em><strong>test</strong></em></p>", "This is a *test*")
|
||||
|
||||
def testParagraphs(self):
|
||||
"""
|
||||
Test extraction to paragraphs
|
||||
"""
|
||||
|
||||
textractor = Textractor(paragraphs=True)
|
||||
|
||||
# Extract text as paragraphs
|
||||
paragraphs = textractor(Utils.PATH + "/article.pdf")
|
||||
|
||||
# Check number of paragraphs is as expected
|
||||
self.assertEqual(len(paragraphs), 11)
|
||||
|
||||
def testSections(self):
|
||||
"""
|
||||
Test extraction to sections
|
||||
"""
|
||||
|
||||
textractor = Textractor(sections=True)
|
||||
|
||||
# Extract as sections
|
||||
sections = textractor(Utils.PATH + "/document.pdf")
|
||||
|
||||
# Check number of sections is as expected
|
||||
self.assertEqual(len(sections), 3)
|
||||
|
||||
def testSentences(self):
|
||||
"""
|
||||
Test extraction to sentences
|
||||
"""
|
||||
|
||||
textractor = Textractor(sentences=True)
|
||||
|
||||
# Extract text as sentences
|
||||
sentences = textractor(Utils.PATH + "/article.pdf")
|
||||
|
||||
# Check number of sentences is as expected
|
||||
self.assertEqual(len(sentences), 17)
|
||||
|
||||
def testSingle(self):
|
||||
"""
|
||||
Test a single extraction with no tokenization of the results
|
||||
"""
|
||||
|
||||
textractor = Textractor()
|
||||
|
||||
# Extract text as a single block
|
||||
text = textractor(Utils.PATH + "/article.pdf")
|
||||
|
||||
# Check length of text is as expected
|
||||
self.assertEqual(len(text), 2471)
|
||||
|
||||
def testTable(self):
|
||||
"""
|
||||
Test table extraction
|
||||
"""
|
||||
|
||||
textractor = Textractor()
|
||||
|
||||
# Extract text as a single block
|
||||
for name in ["document.docx", "spreadsheet.xlsx"]:
|
||||
text = textractor(f"{Utils.PATH}/{name}")
|
||||
|
||||
# Check for table header
|
||||
self.assertTrue("|---|" in text)
|
||||
|
||||
def testTikaFlag(self):
|
||||
"""
|
||||
Test legacy tika flag
|
||||
"""
|
||||
|
||||
textractor = Textractor(tika=True)
|
||||
self.assertIsNotNone(textractor.html)
|
||||
|
||||
textractor = Textractor(tika=False)
|
||||
self.assertIsNone(textractor.html)
|
||||
|
||||
def testTuples(self):
|
||||
"""
|
||||
Test output tuples
|
||||
"""
|
||||
|
||||
# Default text cleaning
|
||||
textractor = Textractor(tuples=True)
|
||||
|
||||
path, text = textractor(Utils.PATH + "/article.pdf")
|
||||
self.assertEqual(path, Utils.PATH + "/article.pdf")
|
||||
self.assertEqual(len(text), 2471)
|
||||
|
||||
def testURL(self):
|
||||
"""
|
||||
Test parsing a remote URL
|
||||
"""
|
||||
|
||||
# Test parsing URLs for each backend
|
||||
for backend in ["docling", "liteparse", "tika"]:
|
||||
textractor = Textractor(backend=backend)
|
||||
text = textractor("https://github.com/neuml/txtai")
|
||||
self.assertTrue("txtai is an all-in-one AI framework" in text)
|
||||
|
||||
def assertMarkdown(self, html, expected):
|
||||
"""
|
||||
Helper method to assert generated markdown is as expected.
|
||||
|
||||
Args:
|
||||
html: input html snippet
|
||||
expected: expected markdown text
|
||||
"""
|
||||
|
||||
textractor = Textractor()
|
||||
self.assertEqual(textractor(f"<html><body>{html}</body></html>"), expected)
|
||||
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
Tokenizer module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Tokenizer
|
||||
|
||||
|
||||
class TestTokenizer(unittest.TestCase):
|
||||
"""
|
||||
Tokenizer tests.
|
||||
"""
|
||||
|
||||
def testAlphanumTokenize(self):
|
||||
"""
|
||||
Test alphanumeric tokenization
|
||||
"""
|
||||
|
||||
# Alphanumeric tokenization through backwards compatible static method
|
||||
self.assertEqual(Tokenizer.tokenize("Y this is a test!"), ["test"])
|
||||
self.assertEqual(Tokenizer.tokenize("abc123 ABC 123"), ["abc123", "abc"])
|
||||
|
||||
def testEmptyTokenize(self):
|
||||
"""
|
||||
Test handling empty and None inputs
|
||||
"""
|
||||
|
||||
# Test that parser can handle empty or None strings
|
||||
self.assertEqual(Tokenizer.tokenize(""), [])
|
||||
self.assertEqual(Tokenizer.tokenize(None), None)
|
||||
|
||||
def testStandardTokenize(self):
|
||||
"""
|
||||
Test standard tokenization
|
||||
"""
|
||||
|
||||
# Default standard tokenizer parameters
|
||||
tokenizer = Tokenizer()
|
||||
|
||||
# Define token tests
|
||||
tests = [
|
||||
("Y this is a test!", ["y", "this", "is", "a", "test"]),
|
||||
("abc123 ABC 123", ["abc123", "abc", "123"]),
|
||||
("Testing hy-phenated words", ["testing", "hy", "phenated", "words"]),
|
||||
("111-111-1111", ["111", "111", "1111"]),
|
||||
("Test.1234", ["test", "1234"]),
|
||||
]
|
||||
|
||||
# Run through tests
|
||||
for test, result in tests:
|
||||
# Unicode Text Segmentation per Unicode Annex #29
|
||||
self.assertEqual(tokenizer(test), result)
|
||||
|
||||
def testNgramTokenize(self):
|
||||
"""
|
||||
Test ngram tokenization
|
||||
"""
|
||||
|
||||
# Standard ngram tokenization
|
||||
tokenizer = Tokenizer(lowercase=True, ngrams=3)
|
||||
result = tokenizer("NGRAM TEST")
|
||||
self.assertIn("ngr", result)
|
||||
|
||||
# Case sensitive ngram tokenization
|
||||
tokenizer = Tokenizer(lowercase=False, ngrams=3)
|
||||
result = tokenizer("NGRAM TEST")
|
||||
self.assertIn("NGR", result)
|
||||
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
URLRetrieve module tests
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import unittest
|
||||
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from threading import Thread
|
||||
from urllib.request import build_opener
|
||||
|
||||
from txtai.pipeline import URLRetrieve
|
||||
from txtai.pipeline.data.urlretrieve import SafeRedirectHandler
|
||||
|
||||
|
||||
class RequestHandler(BaseHTTPRequestHandler):
|
||||
"""
|
||||
Test HTTP handler.
|
||||
"""
|
||||
|
||||
def do_GET(self):
|
||||
"""
|
||||
GET request handler.
|
||||
"""
|
||||
|
||||
if self.path == "/valid":
|
||||
redirect = "https://github.com/neuml/txtai"
|
||||
elif self.path == "/invalid":
|
||||
redirect = "http://127.0.0.1"
|
||||
else:
|
||||
redirect = None
|
||||
|
||||
if redirect:
|
||||
self.send_response(301)
|
||||
self.send_header("Location", redirect)
|
||||
self.end_headers()
|
||||
else:
|
||||
response = "test".encode("utf-8")
|
||||
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", "text/plain")
|
||||
self.send_header("content-length", len(response))
|
||||
self.end_headers()
|
||||
|
||||
self.wfile.write(response)
|
||||
self.wfile.flush()
|
||||
|
||||
|
||||
class TestURLRetrieve(unittest.TestCase):
|
||||
"""
|
||||
URLRetrieve tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create mock http server
|
||||
"""
|
||||
|
||||
cls.httpd = HTTPServer(("127.0.0.1", 8006), RequestHandler)
|
||||
|
||||
server = Thread(target=cls.httpd.serve_forever, daemon=True)
|
||||
server.start()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Shutdown mock http server.
|
||||
"""
|
||||
|
||||
cls.httpd.shutdown()
|
||||
|
||||
def testRedirect(self):
|
||||
"""
|
||||
Test redirects
|
||||
"""
|
||||
|
||||
urlretrieve = URLRetrieve(safeopen=True)
|
||||
|
||||
# Test redirect handler
|
||||
opener = build_opener(SafeRedirectHandler(urlretrieve))
|
||||
|
||||
# Test valid direct
|
||||
with contextlib.closing(opener.open("http://127.0.0.1:8006/valid")) as connection:
|
||||
self.assertTrue("txtai is an all-in-one AI framework" in str(connection.read()))
|
||||
|
||||
# Test invalid redirect
|
||||
with self.assertRaises(IOError):
|
||||
contextlib.closing(opener.open("http://127.0.0.1:8006/invalid"))
|
||||
|
||||
def testRetrieve(self):
|
||||
"""
|
||||
Test retrieval
|
||||
"""
|
||||
|
||||
urlretrieve = URLRetrieve()
|
||||
data = urlretrieve("http://127.0.0.1:8006/data")
|
||||
self.assertEqual(data, b"test")
|
||||
|
||||
def testSafeopen(self):
|
||||
"""
|
||||
Test safeopen checks
|
||||
"""
|
||||
|
||||
urlretrieve = URLRetrieve(safeopen=True)
|
||||
|
||||
# Verify that local ip addresses fail
|
||||
with self.assertRaises(IOError):
|
||||
urlretrieve("http://127.0.0.1")
|
||||
|
||||
with self.assertRaises(IOError):
|
||||
urlretrieve("https://127.0.0.1")
|
||||
|
||||
with self.assertRaises(IOError):
|
||||
urlretrieve("https://[::1]")
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
Caption module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from transformers import AutoModelForImageTextToText, AutoImageProcessor, AutoTokenizer
|
||||
from txtai.pipeline import Caption
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestCaption(unittest.TestCase):
|
||||
"""
|
||||
Caption tests.
|
||||
"""
|
||||
|
||||
def testCaption(self):
|
||||
"""
|
||||
Test captions
|
||||
"""
|
||||
|
||||
caption = Caption()
|
||||
self.assertEqual(caption(Image.open(Utils.PATH + "/books.jpg")), "a book shelf filled with books and a stack of books")
|
||||
|
||||
# Load passing models directly
|
||||
path = "ydshieh/vit-gpt2-coco-en"
|
||||
model = AutoModelForImageTextToText.from_pretrained(path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(path)
|
||||
processor = AutoImageProcessor.from_pretrained(path)
|
||||
|
||||
caption = Caption((model, tokenizer, processor))
|
||||
self.assertEqual(caption(Image.open(Utils.PATH + "/books.jpg")), "a book shelf filled with books and a stack of books")
|
||||
@@ -0,0 +1,74 @@
|
||||
"""
|
||||
ImageHash module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from txtai.pipeline import ImageHash
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestImageHash(unittest.TestCase):
|
||||
"""
|
||||
ImageHash tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Caches an image to hash
|
||||
"""
|
||||
|
||||
cls.image = Image.open(Utils.PATH + "/books.jpg")
|
||||
|
||||
def testArray(self):
|
||||
"""
|
||||
Test numpy return type
|
||||
"""
|
||||
|
||||
ihash = ImageHash(strings=False)
|
||||
self.assertEqual(ihash(self.image).shape, (64,))
|
||||
|
||||
def testAverage(self):
|
||||
"""
|
||||
Test average hash
|
||||
"""
|
||||
|
||||
ihash = ImageHash("average")
|
||||
self.assertIn(ihash(self.image), ["0859dd04bfbfbf00", "0859dd04ffbfbf00"])
|
||||
|
||||
def testColor(self):
|
||||
"""
|
||||
Test color hash
|
||||
"""
|
||||
|
||||
ihash = ImageHash("color")
|
||||
self.assertIn(ihash(self.image), ["1ffffe02000e000c0e0000070000", "1ff8fe03000e00070e0000070000"])
|
||||
|
||||
def testDifference(self):
|
||||
"""
|
||||
Test difference hash
|
||||
"""
|
||||
|
||||
ihash = ImageHash("difference")
|
||||
self.assertEqual(ihash(self.image), "d291996d6969686a")
|
||||
|
||||
def testPerceptual(self):
|
||||
"""
|
||||
Test perceptual hash
|
||||
"""
|
||||
|
||||
ihash = ImageHash("perceptual")
|
||||
self.assertEqual(ihash(self.image), "8be8418577b331b9")
|
||||
|
||||
def testWavelet(self):
|
||||
"""
|
||||
Test wavelet hash
|
||||
"""
|
||||
|
||||
ihash = ImageHash("wavelet")
|
||||
self.assertEqual(ihash(Utils.PATH + "/books.jpg"), "68015d85bfbf3f00")
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Objects module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Objects
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestObjects(unittest.TestCase):
|
||||
"""
|
||||
Object detection tests.
|
||||
"""
|
||||
|
||||
def testClassification(self):
|
||||
"""
|
||||
Test object detection using an image classification model
|
||||
"""
|
||||
|
||||
objects = Objects(classification=True, threshold=0.3)
|
||||
self.assertEqual(objects(Utils.PATH + "/books.jpg")[0][0], "library")
|
||||
|
||||
def testDetection(self):
|
||||
"""
|
||||
Test object detection using an object detection model
|
||||
"""
|
||||
|
||||
objects = Objects()
|
||||
self.assertEqual(objects(Utils.PATH + "/books.jpg")[0][0], "book")
|
||||
|
||||
def testFlatten(self):
|
||||
"""
|
||||
Test object detection using an object detection model, flatten to return only objects
|
||||
"""
|
||||
|
||||
objects = Objects()
|
||||
self.assertEqual(objects(Utils.PATH + "/books.jpg", flatten=True)[0], "book")
|
||||
@@ -0,0 +1,24 @@
|
||||
"""
|
||||
Generator module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Generator
|
||||
|
||||
|
||||
class TestGenerator(unittest.TestCase):
|
||||
"""
|
||||
Sequences tests.
|
||||
"""
|
||||
|
||||
def testGeneration(self):
|
||||
"""
|
||||
Test text pipeline generation
|
||||
"""
|
||||
|
||||
model = Generator("hf-internal-testing/tiny-random-gpt2")
|
||||
start = "Hello, how are"
|
||||
|
||||
# Test that text is generated
|
||||
self.assertIsNotNone(model(start))
|
||||
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
LiteLLM module tests
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import unittest
|
||||
import uuid
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from threading import Thread
|
||||
|
||||
from txtai.pipeline import LLM
|
||||
|
||||
|
||||
class RequestHandler(BaseHTTPRequestHandler):
|
||||
"""
|
||||
Test HTTP handler.
|
||||
"""
|
||||
|
||||
def do_POST(self):
|
||||
"""
|
||||
POST request handler.
|
||||
"""
|
||||
|
||||
# Parse input headers
|
||||
length = int(self.headers["content-length"])
|
||||
data = json.loads(self.rfile.read(length))
|
||||
|
||||
if data.get("stream"):
|
||||
# Mock streaming response
|
||||
content = "application/octet-stream"
|
||||
response = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": str(uuid.uuid4()),
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time() * 1000),
|
||||
"model": "test",
|
||||
"choices": [{"id": 0, "delta": {"content": "blue"}}],
|
||||
}
|
||||
)
|
||||
+ "\n\ndata: [DONE]\n\n"
|
||||
)
|
||||
else:
|
||||
# Mock standard response
|
||||
content = "application/json"
|
||||
response = json.dumps(
|
||||
{
|
||||
"id": str(uuid.uuid4()),
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time() * 1000),
|
||||
"model": "test",
|
||||
"choices": [{"id": 0, "message": {"role": "assistant", "content": "blue"}, "finish_reason": "stop"}],
|
||||
}
|
||||
)
|
||||
|
||||
# Encode response as bytes
|
||||
response = response.encode("utf-8")
|
||||
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", content)
|
||||
self.send_header("content-length", len(response))
|
||||
self.end_headers()
|
||||
|
||||
self.wfile.write(response)
|
||||
self.wfile.flush()
|
||||
|
||||
|
||||
class TestLiteLLM(unittest.TestCase):
|
||||
"""
|
||||
LiteLLM tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create mock http server.
|
||||
"""
|
||||
|
||||
cls.httpd = HTTPServer(("127.0.0.1", 8000), RequestHandler)
|
||||
|
||||
server = Thread(target=cls.httpd.serve_forever, daemon=True)
|
||||
server.start()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Shutdown mock http server.
|
||||
"""
|
||||
|
||||
cls.httpd.shutdown()
|
||||
|
||||
@patch.dict(os.environ, {"OPENAI_API_KEY": "test"})
|
||||
def testGeneration(self):
|
||||
"""
|
||||
Test generation with LiteLLM
|
||||
"""
|
||||
|
||||
# Test model generation with LiteLLM
|
||||
model = LLM("openai/gpt-4o", api_base="http://127.0.0.1:8000")
|
||||
self.assertEqual(model("The sky is"), "blue")
|
||||
|
||||
# Test default role
|
||||
self.assertEqual(model("The sky is", defaultrole="user"), "blue")
|
||||
|
||||
# Test streaming
|
||||
self.assertEqual(" ".join(x for x in model("The sky is", stream=True)), "blue")
|
||||
|
||||
# Test vision
|
||||
self.assertEqual(model.isvision(), False)
|
||||
@@ -0,0 +1,31 @@
|
||||
"""
|
||||
LiteRT module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import LLM
|
||||
|
||||
|
||||
class TestLiteRT(unittest.TestCase):
|
||||
"""
|
||||
LiteRT tests.
|
||||
"""
|
||||
|
||||
def testGeneration(self):
|
||||
"""
|
||||
Test generation with LiteRT
|
||||
"""
|
||||
|
||||
# Test model generation with LiteRT
|
||||
model = LLM("neuml/gemma-4-tiny-random-litert-lm/gemma-4-tiny-random.litertlm", mtp=False, maxlength=25)
|
||||
|
||||
# Test standard
|
||||
self.assertIsNotNone(model("Hello"))
|
||||
|
||||
# Test streaming
|
||||
self.assertIsNotNone(list(model("Hello", stream=True)))
|
||||
|
||||
# Test CPU fallback
|
||||
model = LLM("neuml/gemma-4-tiny-random-litert-lm/gemma-4-tiny-random.litertlm", mtp=True, maxlength=25)
|
||||
self.assertIsNotNone(model("Hello"))
|
||||
@@ -0,0 +1,76 @@
|
||||
"""
|
||||
Llama module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from txtai.pipeline import LLM
|
||||
|
||||
|
||||
class TestLlama(unittest.TestCase):
|
||||
"""
|
||||
llama.cpp tests.
|
||||
"""
|
||||
|
||||
@patch("llama_cpp.Llama")
|
||||
def testContext(self, llama):
|
||||
"""
|
||||
Test n_ctx with llama.cpp
|
||||
"""
|
||||
|
||||
class Llama:
|
||||
"""
|
||||
Mock llama.cpp instance to test invalid context
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
if kwargs.get("n_ctx") == 0 or kwargs.get("n_ctx", 0) >= 10000:
|
||||
raise ValueError("Failed to create context")
|
||||
|
||||
# Save parameters
|
||||
self.params = kwargs
|
||||
|
||||
# Mock llama.cpp instance
|
||||
llama.side_effect = Llama
|
||||
|
||||
# Model to test
|
||||
path = "TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF/tinyllama-1.1b-chat-v0.3.Q2_K.gguf"
|
||||
|
||||
# Test omitting n_ctx falls back to default settings
|
||||
llm = LLM(path)
|
||||
self.assertNotIn("n_ctx", llm.generator.llm.params)
|
||||
|
||||
# Test n_ctx=0 falls back to default settings
|
||||
llm = LLM(path, n_ctx=0)
|
||||
self.assertNotIn("n_ctx", llm.generator.llm.params)
|
||||
|
||||
# Test n_ctx manually set
|
||||
llm = LLM(path, n_ctx=1024)
|
||||
self.assertEqual(llm.generator.llm.params["n_ctx"], 1024)
|
||||
|
||||
# Mock a value for n_ctx that's too big
|
||||
with self.assertRaises(ValueError):
|
||||
llm = LLM(path, n_ctx=10000)
|
||||
|
||||
def testGeneration(self):
|
||||
"""
|
||||
Test generation with llama.cpp
|
||||
"""
|
||||
|
||||
# Test model generation with llama.cpp
|
||||
model = LLM("TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF/tinyllama-1.1b-chat-v0.3.Q2_K.gguf", chat_format="chatml")
|
||||
|
||||
# Test with prompt
|
||||
self.assertEqual(model("2 + 2 = ", maxlength=10, seed=0, stop=["."], defaultrole="prompt")[0], "4")
|
||||
|
||||
# Test with list of messages
|
||||
messages = [{"role": "system", "content": "You are a helpful assistant. You answer math problems."}, {"role": "user", "content": "2+2?"}]
|
||||
self.assertIsNotNone(model(messages, maxlength=10, seed=0, stop=["."]))
|
||||
|
||||
# Test default role
|
||||
self.assertIsNotNone(model("2 + 2 = ", maxlength=10, seed=0, stop=["."], defaultrole="user"))
|
||||
|
||||
# Test streaming
|
||||
self.assertEqual(" ".join(x for x in model("2 + 2 = ", maxlength=10, stream=True, seed=0, stop=["."], defaultrole="prompt"))[0], "4")
|
||||
@@ -0,0 +1,185 @@
|
||||
"""
|
||||
LLM module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from txtai.pipeline import LLM, Generation
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
class TestLLM(unittest.TestCase):
|
||||
"""
|
||||
LLM tests.
|
||||
"""
|
||||
|
||||
def testArguments(self):
|
||||
"""
|
||||
Test pipeline keyword arguments
|
||||
"""
|
||||
|
||||
start = "Hello, how are"
|
||||
|
||||
# Test that text is generated with custom parameters
|
||||
model = LLM("hf-internal-testing/tiny-random-gpt2", task="language-generation", dtype="torch.float32")
|
||||
self.assertIsNotNone(model(start))
|
||||
|
||||
model = LLM("hf-internal-testing/tiny-random-gpt2", task="language-generation", dtype=torch.float32)
|
||||
self.assertIsNotNone(model(start))
|
||||
|
||||
def testBatchSize(self):
|
||||
"""
|
||||
Test batch size
|
||||
"""
|
||||
|
||||
model = LLM("sshleifer/tiny-gpt2")
|
||||
self.assertIsNotNone(model(["Hello, how are"] * 2, batch_size=2))
|
||||
|
||||
def testCustom(self):
|
||||
"""
|
||||
Test custom LLM framework
|
||||
"""
|
||||
|
||||
model = LLM("hf-internal-testing/tiny-random-gpt2", task="language-generation", method="txtai.pipeline.HFGeneration")
|
||||
self.assertIsNotNone(model("Hello, how are"))
|
||||
|
||||
def testCustomNotFound(self):
|
||||
"""
|
||||
Test resolving an unresolvable LLM framework
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
LLM("hf-internal-testing/tiny-random-gpt2", method="notfound.generation")
|
||||
|
||||
def testDefaultRole(self):
|
||||
"""
|
||||
Test default role
|
||||
"""
|
||||
|
||||
model = LLM("hf-internal-testing/tiny-random-LlamaForCausalLM")
|
||||
generator = model.generator
|
||||
|
||||
# Validate that the LLM supports chat messages
|
||||
self.assertEqual(model.ischat(), True)
|
||||
|
||||
messages = [
|
||||
("Hello", list),
|
||||
("\n<|im_start|>Hello<|im_end|>", str),
|
||||
("<|start|>Hello<|end|>", str),
|
||||
("<|start_of_role|>system<|end_of_role|>", str),
|
||||
("[INST]Hello[/INST]", str),
|
||||
]
|
||||
|
||||
for message, expected in messages:
|
||||
# Test auto detection of formats
|
||||
self.assertEqual(type(generator.format([message], "auto")[0]), expected)
|
||||
|
||||
# Test always setting user chat messages
|
||||
self.assertEqual(type(generator.format([message], "user")[0]), list)
|
||||
|
||||
# Test always keeping as prompt text
|
||||
self.assertEqual(type(generator.format([message], "prompt")[0]), str)
|
||||
|
||||
def testExternal(self):
|
||||
"""
|
||||
Test externally loaded model
|
||||
"""
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
|
||||
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
|
||||
|
||||
model = LLM((model, tokenizer), template="{text}")
|
||||
start = "Hello, how are"
|
||||
|
||||
# Test that text is generated
|
||||
self.assertIsNotNone(model(start))
|
||||
|
||||
def testMaxLength(self):
|
||||
"""
|
||||
Test max length
|
||||
"""
|
||||
|
||||
model = LLM("sshleifer/tiny-gpt2")
|
||||
self.assertIsInstance(model("Hello, how are", maxlength=10), str)
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
generation = Generation()
|
||||
self.assertRaises(NotImplementedError, generation.stream, None, None, None, None)
|
||||
|
||||
def testStop(self):
|
||||
"""
|
||||
Test stop strings
|
||||
"""
|
||||
|
||||
model = LLM("sshleifer/tiny-gpt2")
|
||||
self.assertIsNotNone(model("Hello, how are", stop=["you"]))
|
||||
|
||||
def testStream(self):
|
||||
"""
|
||||
Test streaming generation
|
||||
"""
|
||||
|
||||
model = LLM("sshleifer/tiny-gpt2")
|
||||
self.assertIsInstance(" ".join(x for x in model("Hello, how are", stream=True)), str)
|
||||
|
||||
def testStripThink(self):
|
||||
"""
|
||||
Test stripthink parameter
|
||||
"""
|
||||
|
||||
# pylint: disable=W0613
|
||||
def execute1(*args, **kwargs):
|
||||
return ["<think>test</think>you"]
|
||||
|
||||
def execute2(*args, **kwargs):
|
||||
return ["<|channel|>final<|message|> you"]
|
||||
|
||||
model = LLM("hf-internal-testing/tiny-random-LlamaForCausalLM")
|
||||
|
||||
for method in [execute1, execute2]:
|
||||
# Override execute method
|
||||
model.generator.execute = method
|
||||
self.assertEqual(model("Hello, how are", stripthink=True), "you")
|
||||
self.assertEqual(model("Hello, how are", stripthink=False), method()[0])
|
||||
|
||||
def testStripThinkStream(self):
|
||||
"""
|
||||
Test stripthink parameter with streaming output
|
||||
"""
|
||||
|
||||
# pylint: disable=W0613
|
||||
def execute1(*args, **kwargs):
|
||||
yield from "<think>test</think>you"
|
||||
|
||||
def execute2(*args, **kwargs):
|
||||
yield from "<|channel|>final<|message|>you"
|
||||
|
||||
model = LLM("hf-internal-testing/tiny-random-LlamaForCausalLM")
|
||||
|
||||
for method in [execute1, execute2]:
|
||||
# Override execute method
|
||||
model.generator.execute = method
|
||||
self.assertEqual("".join(model("Hello, how are", stripthink=True, stream=True)), "you")
|
||||
self.assertEqual("".join(model("Hello, how are", stripthink=False, stream=True)), "".join(list(method())))
|
||||
|
||||
def testVision(self):
|
||||
"""
|
||||
Test vision LLM
|
||||
"""
|
||||
|
||||
model = LLM("neuml/tiny-random-qwen2vl")
|
||||
result = model(
|
||||
[{"role": "user", "content": [{"type": "text", "text": "What is in this image?"}, {"type": "image", "image": Utils.PATH + "/books.jpg"}]}]
|
||||
)
|
||||
|
||||
self.assertIsNotNone(result)
|
||||
@@ -0,0 +1,71 @@
|
||||
"""
|
||||
OpenCode module tests
|
||||
"""
|
||||
|
||||
import json
|
||||
import unittest
|
||||
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from threading import Thread
|
||||
|
||||
from txtai.pipeline import LLM
|
||||
|
||||
|
||||
class RequestHandler(BaseHTTPRequestHandler):
|
||||
"""
|
||||
Test HTTP handler.
|
||||
"""
|
||||
|
||||
def do_POST(self):
|
||||
"""
|
||||
POST request handler.
|
||||
"""
|
||||
|
||||
# Mock response
|
||||
content = "application/json"
|
||||
response = json.dumps({"id": "0", "parts": [{"type": "text", "text": "blue"}]})
|
||||
|
||||
# Encode response as bytes
|
||||
response = response.encode("utf-8")
|
||||
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", content)
|
||||
self.send_header("content-length", len(response))
|
||||
self.end_headers()
|
||||
|
||||
self.wfile.write(response)
|
||||
self.wfile.flush()
|
||||
|
||||
|
||||
class TestOpenCode(unittest.TestCase):
|
||||
"""
|
||||
OpenCode tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create mock http server.
|
||||
"""
|
||||
|
||||
cls.httpd = HTTPServer(("127.0.0.1", 8005), RequestHandler)
|
||||
|
||||
server = Thread(target=cls.httpd.serve_forever, daemon=True)
|
||||
server.start()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Shutdown mock http server.
|
||||
"""
|
||||
|
||||
cls.httpd.shutdown()
|
||||
|
||||
def testGeneration(self):
|
||||
"""
|
||||
Test generation with OpenCode
|
||||
"""
|
||||
|
||||
# Test model generation with LiteLLM
|
||||
model = LLM("opencode/big-pickle", url="http://127.0.0.1:8005")
|
||||
self.assertEqual(model("The sky is"), "blue")
|
||||
@@ -0,0 +1,225 @@
|
||||
"""
|
||||
RAG module tests
|
||||
"""
|
||||
|
||||
import platform
|
||||
import unittest
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
from txtai.pipeline import Questions, RAG, Similarity
|
||||
|
||||
|
||||
class TestRAG(unittest.TestCase):
|
||||
"""
|
||||
RAG tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create single rag instance.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"Giants hit 3 HRs to down Dodgers",
|
||||
"Giants 5 Dodgers 4 final",
|
||||
"Dodgers drop Game 2 against the Giants, 5-4",
|
||||
"Blue Jays beat Red Sox final score 2-1",
|
||||
"Red Sox lost to the Blue Jays, 2-1",
|
||||
"Blue Jays at Red Sox is over. Score: 2-1",
|
||||
"Phillies win over the Braves, 5-0",
|
||||
"Phillies 5 Braves 0 final",
|
||||
"Final: Braves lose to the Phillies in the series opener, 5-0",
|
||||
"Lightning goaltender pulled, lose to Flyers 4-1",
|
||||
"Flyers 4 Lightning 1 final",
|
||||
"Flyers win 4-1",
|
||||
]
|
||||
|
||||
# Create embeddings model, backed by sentence-transformers & transformers
|
||||
cls.embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
|
||||
# Create rag instance
|
||||
cls.rag = RAG(cls.embeddings, "distilbert-base-cased-distilled-squad")
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Cleanup data.
|
||||
"""
|
||||
|
||||
if cls.embeddings:
|
||||
cls.embeddings.close()
|
||||
|
||||
def testAnswer(self):
|
||||
"""
|
||||
Test qa extraction with an answer
|
||||
"""
|
||||
|
||||
questions = ["What team won the game?", "What was score?"]
|
||||
|
||||
# pylint: disable=C3001
|
||||
execute = lambda query: self.rag([(question, query, question, False) for question in questions], self.data)
|
||||
|
||||
answers = execute("Red Sox - Blue Jays")
|
||||
self.assertEqual("Blue Jays", answers[0][1])
|
||||
self.assertEqual("2-1", answers[1][1])
|
||||
|
||||
# Ad-hoc questions
|
||||
question = "What hockey team won?"
|
||||
|
||||
answers = self.rag([(question, question, question, False)], self.data)
|
||||
self.assertEqual("Flyers", answers[0][1])
|
||||
|
||||
def testEmptyQuery(self):
|
||||
"""
|
||||
Test an empty queries list
|
||||
"""
|
||||
|
||||
self.assertEqual(self.rag.query(None, None), [])
|
||||
|
||||
def testNoAnswer(self):
|
||||
"""
|
||||
Test qa extraction with no answer
|
||||
"""
|
||||
|
||||
question = ""
|
||||
|
||||
answers = self.rag([(question, question, question, False)], self.data)
|
||||
self.assertIsNone(answers[0][1])
|
||||
|
||||
question = "abcdef"
|
||||
answers = self.rag([(question, question, question, False)], self.data)
|
||||
self.assertIsNone(answers[0][1])
|
||||
|
||||
@unittest.skipIf(platform.system() == "Darwin", "Quantized models not supported on macOS")
|
||||
def testQuantize(self):
|
||||
"""
|
||||
Test qa extraction backed by a quantized model
|
||||
"""
|
||||
|
||||
rag = RAG(self.embeddings, "distilbert-base-cased-distilled-squad", True)
|
||||
|
||||
question = "How many home runs?"
|
||||
|
||||
answers = rag([(question, question, question, True)], self.data)
|
||||
self.assertIsNotNone(answers[0][1])
|
||||
|
||||
def testOutputs(self):
|
||||
"""
|
||||
Test output formatting rules
|
||||
"""
|
||||
|
||||
question = "How many home runs?"
|
||||
|
||||
# Test flatten to list of answers
|
||||
rag = RAG(self.embeddings, "distilbert-base-cased-distilled-squad", output="flatten")
|
||||
answers = rag([(question, question, question, True)], self.data)
|
||||
self.assertTrue(answers[0].startswith("Giants hit 3 HRs"))
|
||||
|
||||
# Test reference field
|
||||
rag = RAG(self.embeddings, "distilbert-base-cased-distilled-squad", output="reference")
|
||||
answers = rag([(question, question, question, True)], self.data)
|
||||
self.assertTrue(self.data[answers[0][2]].startswith("Giants hit 3 HRs"))
|
||||
|
||||
def testPrompt(self):
|
||||
"""
|
||||
Test a user prompt with templating
|
||||
"""
|
||||
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
rag = RAG(
|
||||
embeddings,
|
||||
"google/flan-t5-small",
|
||||
template="""
|
||||
Answer the following question and return a number.
|
||||
Question: {question}
|
||||
Context:{context}""",
|
||||
output="flatten",
|
||||
)
|
||||
|
||||
self.assertEqual(rag("How many HRs"), "3")
|
||||
|
||||
def testPromptTemplates(self):
|
||||
"""
|
||||
Test system and user prompt templates
|
||||
"""
|
||||
|
||||
rag = RAG(
|
||||
self.embeddings,
|
||||
"sshleifer/tiny-gpt2",
|
||||
system="You are a friendly assistant",
|
||||
template="""
|
||||
Answer the following question and return a number.
|
||||
Question: {question}
|
||||
Context:{context}""",
|
||||
)
|
||||
|
||||
prompts = rag.prompts(["How many HRs?"], [self.data])[0]
|
||||
self.assertEqual([x["role"] for x in prompts], ["system", "user"])
|
||||
|
||||
def testSearch(self):
|
||||
"""
|
||||
Test qa extraction with an embeddings search for context
|
||||
"""
|
||||
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2", "content": True})
|
||||
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
|
||||
|
||||
rag = RAG(embeddings, "distilbert-base-cased-distilled-squad")
|
||||
|
||||
question = "How many home runs?"
|
||||
|
||||
answers = rag([(question, question, question, True)])
|
||||
self.assertTrue(answers[0][1].startswith("Giants hit 3 HRs"))
|
||||
|
||||
def testSimilarity(self):
|
||||
"""
|
||||
Test qa extraction using a Similarity pipeline to build context
|
||||
"""
|
||||
|
||||
# Create rag instance
|
||||
rag = RAG(Similarity("prajjwal1/bert-medium-mnli"), Questions("distilbert-base-cased-distilled-squad"))
|
||||
|
||||
question = "How many home runs?"
|
||||
|
||||
answers = rag([(question, "HRs", question, True)], self.data)
|
||||
self.assertTrue(answers[0][1].startswith("Giants hit 3 HRs"))
|
||||
|
||||
def testSnippet(self):
|
||||
"""
|
||||
Test qa extraction with a full answer snippet
|
||||
"""
|
||||
|
||||
question = "How many home runs?"
|
||||
|
||||
answers = self.rag([(question, question, question, True)], self.data)
|
||||
self.assertTrue(answers[0][1].startswith("Giants hit 3 HRs"))
|
||||
|
||||
def testSnippetEmpty(self):
|
||||
"""
|
||||
Test snippet method can handle empty parameters
|
||||
"""
|
||||
|
||||
self.assertEqual(self.rag.snippets(["name"], [None], [None], [None]), [("name", None)])
|
||||
|
||||
def testStringInput(self):
|
||||
"""
|
||||
Test with single string input
|
||||
"""
|
||||
|
||||
result = self.rag("How many home runs?", self.data)
|
||||
self.assertEqual(result["answer"], "3")
|
||||
|
||||
def testTasks(self):
|
||||
"""
|
||||
Test loading models with task parameter
|
||||
"""
|
||||
|
||||
for task, model in [
|
||||
("language-generation", "hf-internal-testing/tiny-random-gpt2"),
|
||||
("sequence-sequence", "hf-internal-testing/tiny-random-t5"),
|
||||
]:
|
||||
rag = RAG(self.embeddings, model, task=task)
|
||||
self.assertIsNotNone(rag)
|
||||
@@ -0,0 +1,21 @@
|
||||
"""
|
||||
Sequences module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Sequences
|
||||
|
||||
|
||||
class TestSequences(unittest.TestCase):
|
||||
"""
|
||||
Sequences tests.
|
||||
"""
|
||||
|
||||
def testGeneration(self):
|
||||
"""
|
||||
Test text2text pipeline generation
|
||||
"""
|
||||
|
||||
model = Sequences("t5-small")
|
||||
self.assertEqual(model("Testing the model", prefix="translate English to German: "), "Das Modell zu testen")
|
||||
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
Entity module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Entity
|
||||
|
||||
|
||||
class TestEntity(unittest.TestCase):
|
||||
"""
|
||||
Entity tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create entity instance.
|
||||
"""
|
||||
|
||||
cls.entity = Entity("dslim/bert-base-NER")
|
||||
|
||||
def testEntity(self):
|
||||
"""
|
||||
Test entity
|
||||
"""
|
||||
|
||||
# Run entity extraction
|
||||
entities = self.entity("Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg")
|
||||
self.assertEqual([e[0] for e in entities], ["Canada", "Manhattan"])
|
||||
|
||||
def testEntityFlatten(self):
|
||||
"""
|
||||
Test entity with flattened output
|
||||
"""
|
||||
|
||||
# Test flatten
|
||||
entities = self.entity("Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg", flatten=True)
|
||||
self.assertEqual(entities, ["Canada", "Manhattan"])
|
||||
|
||||
# Test flatten with join
|
||||
entities = self.entity(
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg", flatten=True, join=True
|
||||
)
|
||||
self.assertEqual(entities, "Canada Manhattan")
|
||||
|
||||
def testEntityTypes(self):
|
||||
"""
|
||||
Test entity type filtering
|
||||
"""
|
||||
|
||||
# Run entity extraction
|
||||
entities = self.entity("Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg", labels=["PER"])
|
||||
self.assertFalse(entities)
|
||||
|
||||
def testGliner(self):
|
||||
"""
|
||||
Test entity pipeline with a GLiNER model
|
||||
"""
|
||||
|
||||
entity = Entity("neuml/gliner-bert-tiny")
|
||||
entities = entity("My name is John Smith.", flatten=True)
|
||||
self.assertEqual(entities, ["John Smith"])
|
||||
@@ -0,0 +1,85 @@
|
||||
"""
|
||||
Labels module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Labels
|
||||
|
||||
|
||||
class TestLabels(unittest.TestCase):
|
||||
"""
|
||||
Labels tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create single labels instance.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
cls.labels = Labels("prajjwal1/bert-medium-mnli")
|
||||
|
||||
def testLabel(self):
|
||||
"""
|
||||
Test labels with single text input
|
||||
"""
|
||||
|
||||
self.assertEqual(self.labels("This is the best sentence ever", ["positive", "negative"])[0][0], 0)
|
||||
|
||||
def testLabelFlatten(self):
|
||||
"""
|
||||
Test labels with single text input, flattened to top text labels
|
||||
"""
|
||||
|
||||
self.assertEqual(self.labels("This is the best sentence ever", ["positive", "negative"], flatten=True)[0], "positive")
|
||||
|
||||
def testLabelBatch(self):
|
||||
"""
|
||||
Test labels with multiple text inputs
|
||||
"""
|
||||
|
||||
results = [l[0][0] for l in self.labels(["This is the best sentence ever", "This is terrible"], ["positive", "negative"])]
|
||||
self.assertEqual(results, [0, 1])
|
||||
|
||||
def testLabelBatchFlatten(self):
|
||||
"""
|
||||
Test labels with multiple text inputs, flattened to top text labels
|
||||
"""
|
||||
|
||||
results = [l[0] for l in self.labels(["This is the best sentence ever", "This is terrible"], ["positive", "negative"], flatten=True)]
|
||||
self.assertEqual(results, ["positive", "negative"])
|
||||
|
||||
def testLabelFixed(self):
|
||||
"""
|
||||
Test labels with a fixed label text classification model
|
||||
"""
|
||||
|
||||
labels = Labels(dynamic=False)
|
||||
|
||||
# Get index of "POSITIVE" label
|
||||
index = labels.labels().index("POSITIVE")
|
||||
|
||||
# Verify results
|
||||
self.assertEqual(labels("This is the best sentence ever")[0][0], index)
|
||||
self.assertEqual(labels("This is the best sentence ever", multilabel=True)[0][0], index)
|
||||
|
||||
def testLabelFixedFlatten(self):
|
||||
"""
|
||||
Test labels with a fixed label text classification model, flattened to top text labels
|
||||
"""
|
||||
|
||||
labels = Labels(dynamic=False)
|
||||
|
||||
# Verify results
|
||||
self.assertEqual(labels("This is the best sentence ever", flatten=True)[0], "POSITIVE")
|
||||
self.assertEqual(labels("This is the best sentence ever", multilabel=True, flatten=True)[0], "POSITIVE")
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
Reranker module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai import Embeddings
|
||||
from txtai.pipeline import Reranker, Similarity
|
||||
|
||||
|
||||
class TestReranker(unittest.TestCase):
|
||||
"""
|
||||
Reranker tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create single labels instance.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
def testRanker(self):
|
||||
"""
|
||||
Test re-ranking pipeline
|
||||
"""
|
||||
|
||||
embeddings = Embeddings(content=True)
|
||||
embeddings.index(self.data)
|
||||
|
||||
similarity = Similarity("neuml/colbert-bert-tiny", lateencode=True)
|
||||
|
||||
ranker = Reranker(embeddings, similarity)
|
||||
self.assertEqual(ranker("lottery winner")[0]["id"], "4")
|
||||
@@ -0,0 +1,105 @@
|
||||
"""
|
||||
Similarity module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Similarity
|
||||
|
||||
|
||||
class TestSimilarity(unittest.TestCase):
|
||||
"""
|
||||
Similarity tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create single labels instance.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
cls.similarity = Similarity("prajjwal1/bert-medium-mnli")
|
||||
|
||||
def testCrossEncoder(self):
|
||||
"""
|
||||
Test cross-encoder similarity model
|
||||
"""
|
||||
|
||||
similarity = Similarity("cross-encoder/ms-marco-MiniLM-L-2-v2", crossencode=True)
|
||||
uid = similarity("Who won the lottery?", self.data)[0][0]
|
||||
self.assertEqual(self.data[uid], self.data[4])
|
||||
|
||||
def testCrossEncoderBatch(self):
|
||||
"""
|
||||
Test cross-encoder similarity model with multiple inputs
|
||||
"""
|
||||
|
||||
similarity = Similarity("cross-encoder/ms-marco-MiniLM-L-2-v2", crossencode=True)
|
||||
results = [r[0][0] for r in similarity(["Who won the lottery?", "Where did an iceberg collapse?"], self.data)]
|
||||
self.assertEqual(results, [4, 1])
|
||||
|
||||
def testLateEncoder(self):
|
||||
"""
|
||||
Test late-encoder similarity model
|
||||
"""
|
||||
|
||||
similarity = Similarity("neuml/pylate-bert-tiny", lateencode=True)
|
||||
uid = similarity("Who won the lottery?", self.data)[0][0]
|
||||
self.assertEqual(self.data[uid], self.data[4])
|
||||
|
||||
# Test encode method
|
||||
# pylint: disable=E1101
|
||||
self.assertEqual(similarity.encode(["Who won the lottery?"], "data").shape, (1, 8, 128))
|
||||
|
||||
def testLateEncoderBatch(self):
|
||||
"""
|
||||
Test late-encoder similarity model with multiple inputs
|
||||
"""
|
||||
|
||||
similarity = Similarity("neuml/colbert-bert-tiny", lateencode=True)
|
||||
results = [r[0][0] for r in similarity(["Who won the lottery?", "Where did an iceberg collapse?"], self.data)]
|
||||
self.assertEqual(results, [4, 1])
|
||||
|
||||
def testSimilarity(self):
|
||||
"""
|
||||
Test similarity with single query
|
||||
"""
|
||||
|
||||
uid = self.similarity("feel good story", self.data)[0][0]
|
||||
self.assertEqual(self.data[uid], self.data[4])
|
||||
|
||||
def testSimilarityBatch(self):
|
||||
"""
|
||||
Test similarity with multiple queries
|
||||
"""
|
||||
|
||||
results = [r[0][0] for r in self.similarity(["feel good story", "climate change"], self.data)]
|
||||
self.assertEqual(results, [4, 1])
|
||||
|
||||
def testSimilarityFixed(self):
|
||||
"""
|
||||
Test similarity with a fixed label text classification model
|
||||
"""
|
||||
|
||||
similarity = Similarity(dynamic=False)
|
||||
|
||||
# Test with query as label text and label id
|
||||
self.assertLessEqual(similarity("negative", ["This is the best sentence ever"])[0][1], 0.1)
|
||||
self.assertLessEqual(similarity("0", ["This is the best sentence ever"])[0][1], 0.1)
|
||||
|
||||
def testSimilarityLong(self):
|
||||
"""
|
||||
Test similarity with long text
|
||||
"""
|
||||
|
||||
uid = self.similarity("other", ["Very long text " * 1000, "other text"])[0][0]
|
||||
self.assertEqual(uid, 1)
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Summary module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.pipeline import Summary
|
||||
|
||||
|
||||
class TestSummary(unittest.TestCase):
|
||||
"""
|
||||
Summary tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create single summary instance.
|
||||
"""
|
||||
|
||||
cls.text = (
|
||||
"Search is the base of many applications. Once data starts to pile up, users want to be able to find it. It's the foundation "
|
||||
"of the internet and an ever-growing challenge that is never solved or done. The field of Natural Language Processing (NLP) is "
|
||||
"rapidly evolving with a number of new developments. Large-scale general language models are an exciting new capability "
|
||||
"allowing us to add amazing functionality quickly with limited compute and people. Innovation continues with new models "
|
||||
"and advancements coming in at what seems a weekly basis. This article introduces txtai, an AI-powered search engine "
|
||||
"that enables Natural Language Understanding (NLU) based search in any application."
|
||||
)
|
||||
|
||||
cls.summary = Summary("t5-small")
|
||||
|
||||
def testSummary(self):
|
||||
"""
|
||||
Test summarization of text
|
||||
"""
|
||||
|
||||
self.assertEqual(self.summary(self.text, minlength=15, maxlength=15), "the field of natural language processing (NLP) is rapidly evolving")
|
||||
|
||||
def testSummaryBatch(self):
|
||||
"""
|
||||
Test batch summarization of text
|
||||
"""
|
||||
|
||||
summaries = self.summary([self.text, self.text], maxlength=15)
|
||||
self.assertEqual(len(summaries), 2)
|
||||
|
||||
def testSummaryNoLength(self):
|
||||
"""
|
||||
Test summary with no max length set
|
||||
"""
|
||||
|
||||
self.assertEqual(
|
||||
self.summary(self.text + self.text),
|
||||
"search is the base of many applications. Once data starts to pile up, users want to be able to find it. "
|
||||
+ "Large-scale general language models are an exciting new capability allowing us to add amazing functionality quickly "
|
||||
+ "with limited compute and people.",
|
||||
)
|
||||
|
||||
def testSummaryShort(self):
|
||||
"""
|
||||
Test that summarization is skipped
|
||||
"""
|
||||
|
||||
self.assertEqual(self.summary("Text", maxlength=15), "Text")
|
||||
@@ -0,0 +1,175 @@
|
||||
"""
|
||||
Translation module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from txtai.pipeline import Translation
|
||||
|
||||
|
||||
class TestTranslation(unittest.TestCase):
|
||||
"""
|
||||
Translation tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create single translation instance.
|
||||
"""
|
||||
|
||||
cls.translate = Translation()
|
||||
|
||||
# Preload list of models. Handle HF Hub errors.
|
||||
complete, wait = False, 1
|
||||
while not complete:
|
||||
try:
|
||||
cls.translate.lookup("en", "es")
|
||||
complete = True
|
||||
except requests.exceptions.HTTPError:
|
||||
# Exponential backoff
|
||||
time.sleep(wait)
|
||||
|
||||
# Wait up to 16 seconds
|
||||
wait = min(wait * 2, 16)
|
||||
|
||||
def testDetect(self):
|
||||
"""
|
||||
Test language detection
|
||||
"""
|
||||
|
||||
test = ["This is a test language detection."]
|
||||
language = self.translate.detect(test)
|
||||
|
||||
self.assertListEqual(language, ["en"])
|
||||
|
||||
def testDetectWithCustomFunc(self):
|
||||
"""
|
||||
Test language detection with custom function
|
||||
"""
|
||||
|
||||
def dummy_func(text):
|
||||
return ["en" for x in text]
|
||||
|
||||
translate = Translation(langdetect=dummy_func)
|
||||
|
||||
test = ["This is a test language detection."]
|
||||
language = translate.detect(test)
|
||||
|
||||
self.assertListEqual(language, ["en"])
|
||||
|
||||
def testLongTranslation(self):
|
||||
"""
|
||||
Test a translation longer than max tokenization length
|
||||
"""
|
||||
|
||||
text = "This is a test translation to Spanish. " * 100
|
||||
translation = self.translate(text, "es")
|
||||
|
||||
# Validate translation text
|
||||
self.assertIsNotNone(translation)
|
||||
|
||||
def testM2M100Translation(self):
|
||||
"""
|
||||
Test a translation using M2M100 models
|
||||
"""
|
||||
|
||||
text = self.translate("This is a test translation to Croatian", "hr")
|
||||
|
||||
# Validate translation text
|
||||
self.assertEqual(text, "Ovo je testni prijevod na hrvatski")
|
||||
|
||||
def testMarianTranslation(self):
|
||||
"""
|
||||
Test a translation using Marian models
|
||||
"""
|
||||
|
||||
text = "This is a test translation into Spanish"
|
||||
translation = self.translate(text, "es")
|
||||
|
||||
# Validate translation text
|
||||
self.assertEqual(translation, "Esta es una traducción de prueba al español")
|
||||
|
||||
# Validate translation back
|
||||
translation = self.translate(translation, "en")
|
||||
self.assertEqual(translation, text)
|
||||
|
||||
def testNoLang(self):
|
||||
"""
|
||||
Test no matching language id
|
||||
"""
|
||||
|
||||
self.assertIsNone(self.translate.langid([], "zz"))
|
||||
|
||||
def testNoModel(self):
|
||||
"""
|
||||
Test no known available model found
|
||||
"""
|
||||
|
||||
self.assertEqual(self.translate.modelpath("zz", "en"), "Helsinki-NLP/opus-mt-mul-en")
|
||||
|
||||
def testNoTranslation(self):
|
||||
"""
|
||||
Test translation skipped when text already in destination language
|
||||
"""
|
||||
|
||||
text = "This is a test translation to English"
|
||||
translation = self.translate(text, "en")
|
||||
|
||||
# Validate no translation
|
||||
self.assertEqual(text, translation)
|
||||
|
||||
def testShowmodelsChunked(self):
|
||||
"""
|
||||
Test a long translation with showmodels flag. When text is chunked
|
||||
by the tokenizer, results should still be properly concatenated as
|
||||
a 3-tuple (translation, language, model) rather than a malformed tuple.
|
||||
"""
|
||||
|
||||
text = "This is a test translation to Spanish. " * 100
|
||||
result = self.translate(text, "es", showmodels=True)
|
||||
|
||||
# Result should be a tuple of exactly 3 elements
|
||||
self.assertIsInstance(result, tuple)
|
||||
self.assertEqual(len(result), 3)
|
||||
|
||||
translation, language, modelpath = result
|
||||
|
||||
# Translation should be a single string, not a nested tuple
|
||||
self.assertIsInstance(translation, str)
|
||||
self.assertIsNotNone(translation)
|
||||
self.assertGreater(len(translation), 0)
|
||||
|
||||
# Language and model should be valid strings
|
||||
self.assertEqual(language, "en")
|
||||
self.assertIsInstance(modelpath, str)
|
||||
|
||||
def testTranslationWithShowmodels(self):
|
||||
"""
|
||||
Test a translation using Marian models and showmodels flag to return
|
||||
model and language.
|
||||
"""
|
||||
|
||||
text = "This is a test translation into Spanish"
|
||||
result = self.translate(text, "es", showmodels=True)
|
||||
|
||||
translation, language, modelpath = result
|
||||
# Validate translation text
|
||||
self.assertEqual(translation, "Esta es una traducción de prueba al español")
|
||||
# Validate detected language
|
||||
self.assertEqual(language, "en")
|
||||
# Validate model
|
||||
self.assertEqual(modelpath, "Helsinki-NLP/opus-mt-en-es")
|
||||
|
||||
# Validate translation back
|
||||
result = self.translate(translation, "en", showmodels=True)
|
||||
|
||||
translation, language, modelpath = result
|
||||
self.assertEqual(translation, text)
|
||||
# Validate detected language
|
||||
self.assertEqual(language, "es")
|
||||
# Validate model
|
||||
self.assertEqual(modelpath, "Helsinki-NLP/opus-mt-es-en")
|
||||
@@ -0,0 +1,161 @@
|
||||
"""
|
||||
ONNX module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||||
from sklearn.linear_model import LogisticRegression
|
||||
from sklearn.pipeline import Pipeline
|
||||
|
||||
from txtai.embeddings import Embeddings
|
||||
from txtai.models import OnnxModel
|
||||
from txtai.pipeline import HFOnnx, HFTrainer, Labels, MLOnnx, Questions
|
||||
|
||||
|
||||
class TestOnnx(unittest.TestCase):
|
||||
"""
|
||||
ONNX tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create default datasets.
|
||||
"""
|
||||
|
||||
cls.data = [{"text": "Dogs", "label": 0}, {"text": "dog", "label": 0}, {"text": "Cats", "label": 1}, {"text": "cat", "label": 1}] * 100
|
||||
|
||||
def testDefault(self):
|
||||
"""
|
||||
Test exporting an ONNX model with default parameters
|
||||
"""
|
||||
|
||||
# Export model to ONNX, use default parameters
|
||||
onnx = HFOnnx()
|
||||
model = onnx("google/bert_uncased_L-2_H-128_A-2")
|
||||
|
||||
# Validate model has data
|
||||
self.assertGreater(len(model), 0)
|
||||
|
||||
# Validate model device properly works
|
||||
self.assertEqual(OnnxModel(model).device, -1)
|
||||
|
||||
def testClassification(self):
|
||||
"""
|
||||
Test exporting a classification model to ONNX and running inference
|
||||
"""
|
||||
|
||||
path = "google/bert_uncased_L-2_H-128_A-2"
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer(path, self.data)
|
||||
|
||||
# Output file path
|
||||
output = os.path.join(tempfile.gettempdir(), "onnx")
|
||||
|
||||
# Export model to ONNX
|
||||
onnx = HFOnnx()
|
||||
model = onnx((model, tokenizer), "text-classification", output, True)
|
||||
|
||||
# Test classification
|
||||
labels = Labels((model, path), dynamic=False)
|
||||
self.assertEqual(labels("cat")[0][0], 1)
|
||||
|
||||
@patch("onnxruntime.get_available_providers")
|
||||
@patch("torch.cuda.is_available")
|
||||
def testPooling(self, cuda, providers):
|
||||
"""
|
||||
Test exporting a pooling model to ONNX and running inference
|
||||
"""
|
||||
|
||||
path = "sentence-transformers/paraphrase-MiniLM-L3-v2"
|
||||
|
||||
# Export model to ONNX
|
||||
onnx = HFOnnx()
|
||||
model = onnx(path, "pooling", quantize=True)
|
||||
|
||||
# Test no CUDA and onnxruntime installed
|
||||
cuda.return_value = False
|
||||
providers.return_value = ["CPUExecutionProvider"]
|
||||
|
||||
embeddings = Embeddings({"path": model, "tokenizer": path})
|
||||
self.assertEqual(embeddings.similarity("animal", ["dog", "book", "rug"])[0][0], 0)
|
||||
|
||||
# Test no CUDA and onnxruntime-gpu installed
|
||||
cuda.return_value = False
|
||||
providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
|
||||
embeddings = Embeddings({"path": model, "tokenizer": path})
|
||||
self.assertIsNotNone(embeddings)
|
||||
|
||||
# Test CUDA and only onnxruntime installed
|
||||
cuda.return_value = True
|
||||
providers.return_value = ["CPUExecutionProvider"]
|
||||
|
||||
embeddings = Embeddings({"path": model, "tokenizer": path})
|
||||
self.assertIsNotNone(embeddings)
|
||||
|
||||
# Test CUDA and onnxruntime-gpu installed
|
||||
cuda.return_value = True
|
||||
providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
|
||||
embeddings = Embeddings({"path": model, "tokenizer": path})
|
||||
self.assertIsNotNone(embeddings)
|
||||
|
||||
def testQA(self):
|
||||
"""
|
||||
Test exporting a QA model to ONNX and running inference
|
||||
"""
|
||||
|
||||
path = "distilbert-base-cased-distilled-squad"
|
||||
|
||||
# Export model to ONNX
|
||||
onnx = HFOnnx()
|
||||
model = onnx(path, "question-answering")
|
||||
|
||||
questions = Questions((model, path))
|
||||
self.assertEqual(questions(["What is the price?"], ["The price is $30"])[0], "$30")
|
||||
|
||||
def testScikit(self):
|
||||
"""
|
||||
Test exporting a scikit-learn model to ONNX and running inference
|
||||
"""
|
||||
|
||||
# pylint: disable=W0613
|
||||
def tokenizer(inputs, **kwargs):
|
||||
if isinstance(inputs, str):
|
||||
inputs = [inputs]
|
||||
|
||||
return {"input_ids": [[x] for x in inputs]}
|
||||
|
||||
# Train a scikit-learn model
|
||||
model = Pipeline([("tfidf", TfidfVectorizer()), ("lr", LogisticRegression())])
|
||||
model.fit([x["text"] for x in self.data], [x["label"] for x in self.data])
|
||||
|
||||
# Export model to ONNX
|
||||
onnx = MLOnnx()
|
||||
model = onnx(model)
|
||||
|
||||
# Test classification
|
||||
labels = Labels((model, tokenizer), dynamic=False)
|
||||
self.assertEqual(labels("cat")[0][0], 1)
|
||||
|
||||
def testZeroShot(self):
|
||||
"""
|
||||
Test exporting a zero shot classification model to ONNX and running inference
|
||||
"""
|
||||
|
||||
path = "prajjwal1/bert-medium-mnli"
|
||||
|
||||
# Export model to ONNX
|
||||
onnx = HFOnnx()
|
||||
model = onnx(path, "zero-shot-classification", quantize=True)
|
||||
|
||||
# Test zero shot classification
|
||||
labels = Labels((model, path))
|
||||
self.assertEqual(labels("That is great news", ["negative", "positive"])[0][0], 1)
|
||||
@@ -0,0 +1,35 @@
|
||||
"""
|
||||
Quantization module tests
|
||||
"""
|
||||
|
||||
import platform
|
||||
import unittest
|
||||
|
||||
from transformers import AutoModel
|
||||
|
||||
from txtai.pipeline import HFModel, HFPipeline
|
||||
|
||||
|
||||
class TestQuantization(unittest.TestCase):
|
||||
"""
|
||||
Quantization tests.
|
||||
"""
|
||||
|
||||
@unittest.skipIf(platform.system() == "Darwin", "Quantized models not supported on macOS")
|
||||
def testModel(self):
|
||||
"""
|
||||
Test quantizing a model through HFModel.
|
||||
"""
|
||||
|
||||
model = HFModel(quantize=True, gpu=False)
|
||||
model = model.prepare(AutoModel.from_pretrained("google/bert_uncased_L-2_H-128_A-2"))
|
||||
self.assertIsNotNone(model)
|
||||
|
||||
@unittest.skipIf(platform.system() == "Darwin", "Quantized models not supported on macOS")
|
||||
def testPipeline(self):
|
||||
"""
|
||||
Test quantizing a model through HFPipeline.
|
||||
"""
|
||||
|
||||
pipeline = HFPipeline("text-classification", "google/bert_uncased_L-2_H-128_A-2", True, False)
|
||||
self.assertIsNotNone(pipeline)
|
||||
@@ -0,0 +1,335 @@
|
||||
"""
|
||||
Trainer module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
|
||||
from transformers import AutoTokenizer, AutoModelForSequenceClassification
|
||||
|
||||
from txtai.data import Data
|
||||
from txtai.pipeline import HFTrainer, Labels, Questions, Sequences
|
||||
|
||||
|
||||
class TestTrainer(unittest.TestCase):
|
||||
"""
|
||||
Trainer tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create default datasets.
|
||||
"""
|
||||
|
||||
cls.data = [{"text": "Dogs", "label": 0}, {"text": "dog", "label": 0}, {"text": "Cats", "label": 1}, {"text": "cat", "label": 1}] * 100
|
||||
|
||||
def testBasic(self):
|
||||
"""
|
||||
Test training a model with basic parameters
|
||||
"""
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)
|
||||
|
||||
labels = Labels((model, tokenizer), dynamic=False)
|
||||
self.assertEqual(labels("cat")[0][0], 1)
|
||||
|
||||
def testCLM(self):
|
||||
"""
|
||||
Test training a model with causal language modeling
|
||||
"""
|
||||
|
||||
trainer = HFTrainer()
|
||||
|
||||
# Test default parameters
|
||||
model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation")
|
||||
self.assertIsNotNone(model)
|
||||
|
||||
# Test pack merging
|
||||
model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge="pack")
|
||||
self.assertIsNotNone(model)
|
||||
|
||||
# Test no merging
|
||||
model, _ = trainer("hf-internal-testing/tiny-random-gpt2", self.data, maxlength=16, task="language-generation", merge=None)
|
||||
self.assertIsNotNone(model)
|
||||
|
||||
def testCustom(self):
|
||||
"""
|
||||
Test training a model with custom parameters
|
||||
"""
|
||||
|
||||
# pylint: disable=E1120
|
||||
model = AutoModelForSequenceClassification.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
|
||||
tokenizer = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer(
|
||||
(model, tokenizer),
|
||||
self.data,
|
||||
self.data,
|
||||
columns=("text", "label"),
|
||||
do_eval=True,
|
||||
output_dir=os.path.join(tempfile.gettempdir(), "trainer"),
|
||||
)
|
||||
|
||||
labels = Labels((model, tokenizer), dynamic=False)
|
||||
self.assertEqual(labels("cat")[0][0], 1)
|
||||
|
||||
def testDataFrame(self):
|
||||
"""
|
||||
Test training a model with a mock pandas DataFrame
|
||||
"""
|
||||
|
||||
class TestDataFrame:
|
||||
"""
|
||||
Test DataFrame
|
||||
"""
|
||||
|
||||
def __init__(self, data):
|
||||
# Get list of columns
|
||||
self.columns = list(data[0].keys())
|
||||
|
||||
# Build columnar data view
|
||||
self.data = {}
|
||||
for column in self.columns:
|
||||
self.data[column] = Values([row[column] for row in data])
|
||||
|
||||
def __getitem__(self, column):
|
||||
return self.data[column]
|
||||
|
||||
class Values:
|
||||
"""
|
||||
Test values list
|
||||
"""
|
||||
|
||||
def __init__(self, values):
|
||||
self.values = list(values)
|
||||
|
||||
def __getitem__(self, index):
|
||||
return self.values[index]
|
||||
|
||||
def unique(self):
|
||||
"""
|
||||
Returns a list of unique values.
|
||||
|
||||
Returns:
|
||||
unique list of values
|
||||
"""
|
||||
|
||||
return set(self.values)
|
||||
|
||||
# Mock DataFrame
|
||||
df = TestDataFrame(self.data)
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", df)
|
||||
|
||||
labels = Labels((model, tokenizer), dynamic=False)
|
||||
self.assertEqual(labels("cat")[0][0], 1)
|
||||
|
||||
def testDataset(self):
|
||||
"""
|
||||
Test training a model with a mock Hugging Face Dataset
|
||||
"""
|
||||
|
||||
class TestDataset(torch.utils.data.Dataset):
|
||||
"""
|
||||
Test Dataset
|
||||
"""
|
||||
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
self.unique = lambda _: [0, 1]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, index):
|
||||
return self.data[index]
|
||||
|
||||
def column_names(self):
|
||||
"""
|
||||
Returns column names for this dataset
|
||||
|
||||
Returns:
|
||||
list of columns
|
||||
"""
|
||||
|
||||
return ["text", "label"]
|
||||
|
||||
# pylint: disable=W0613
|
||||
def map(self, fn, batched, num_proc, remove_columns):
|
||||
"""
|
||||
Map each dataset row using fn.
|
||||
|
||||
Args:
|
||||
fn: function
|
||||
batched: batch records
|
||||
|
||||
Returns:
|
||||
updated Dataset
|
||||
"""
|
||||
|
||||
self.data = [fn(x) for x in self.data]
|
||||
return self
|
||||
|
||||
ds = TestDataset(self.data)
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", ds)
|
||||
|
||||
labels = Labels((model, tokenizer), dynamic=False)
|
||||
self.assertEqual(labels("cat")[0][0], 1)
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test an empty training data object
|
||||
"""
|
||||
|
||||
self.assertIsNone(Data(None, None, None).process(None))
|
||||
|
||||
def testKD(self):
|
||||
"""
|
||||
Test knowledge distillation
|
||||
"""
|
||||
|
||||
# Base model
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data)
|
||||
|
||||
# Train with knowledge distillation
|
||||
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", self.data, teacher=(model, tokenizer))
|
||||
|
||||
labels = Labels((model, tokenizer), dynamic=False)
|
||||
self.assertEqual(labels("cat")[0][0], 1)
|
||||
|
||||
def testMLM(self):
|
||||
"""
|
||||
Test training a model with masked language modeling.
|
||||
"""
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, _ = trainer("hf-internal-testing/tiny-random-bert", self.data, task="language-modeling")
|
||||
|
||||
# Test model completed successfully
|
||||
self.assertIsNotNone(model)
|
||||
|
||||
def testMultiLabel(self):
|
||||
"""
|
||||
Test training model with labels provided as a list
|
||||
"""
|
||||
|
||||
data = []
|
||||
for x in self.data:
|
||||
data.append({"text": x["text"], "label": [0.0, 1.0] if x["label"] else [1.0, 0.0]})
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)
|
||||
|
||||
labels = Labels((model, tokenizer), dynamic=False)
|
||||
self.assertEqual(labels("cat")[0][0], 1)
|
||||
|
||||
def testPEFT(self):
|
||||
"""
|
||||
Test training a model with causal language modeling and PEFT
|
||||
"""
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, _ = trainer(
|
||||
"hf-internal-testing/tiny-random-gpt2",
|
||||
self.data,
|
||||
maxlength=16,
|
||||
task="language-generation",
|
||||
quantize=True,
|
||||
lora=True,
|
||||
)
|
||||
|
||||
# Test model completed successfully
|
||||
self.assertIsNotNone(model)
|
||||
|
||||
def testQA(self):
|
||||
"""
|
||||
Test training a QA model
|
||||
"""
|
||||
|
||||
# Training data
|
||||
data = [
|
||||
{"question": "What ingredient?", "context": "1 can whole tomatoes", "answers": "tomatoes"},
|
||||
{"question": "What ingredient?", "context": "Crush 1 tomato", "answers": "tomato"},
|
||||
{"question": "What ingredient?", "context": "1 yellow onion", "answers": "onion"},
|
||||
{"question": "What ingredient?", "context": "Unwrap 2 red onions", "answers": "onions"},
|
||||
{"question": "What ingredient?", "context": "1 red pepper", "answers": "pepper"},
|
||||
{"question": "What ingredient?", "context": "Clean 3 red peppers", "answers": "peppers"},
|
||||
{"question": "What ingredient?", "context": "1 clove garlic", "answers": "garlic"},
|
||||
{"question": "What ingredient?", "context": "Unwrap 3 cloves of garlic", "answers": "garlic"},
|
||||
{"question": "What ingredient?", "context": "3 pieces of ginger", "answers": "ginger"},
|
||||
{"question": "What ingredient?", "context": "Peel 1 orange", "answers": "orange"},
|
||||
{"question": "What ingredient?", "context": "1/2 lb beef", "answers": "beef"},
|
||||
{"question": "What ingredient?", "context": "Roast 3 lbs of beef", "answers": "beef"},
|
||||
{"question": "What ingredient?", "context": "1 pack of chicken", "answers": "chicken"},
|
||||
{"question": "What ingredient?", "context": "Forest through the trees", "answers": None},
|
||||
]
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data, data, task="question-answering", num_train_epochs=40)
|
||||
|
||||
questions = Questions((model, tokenizer), gpu=True)
|
||||
self.assertTrue("onion" in questions(["What ingredient?"], ["Peel 1 onion"])[0])
|
||||
|
||||
def testRegression(self):
|
||||
"""
|
||||
Test training a model with a regression (continuous) output
|
||||
"""
|
||||
|
||||
data = []
|
||||
for x in self.data:
|
||||
data.append({"text": x["text"], "label": x["label"] + 0.1})
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)
|
||||
|
||||
labels = Labels((model, tokenizer), dynamic=False)
|
||||
|
||||
# Regression tasks return a single entry with the regression output
|
||||
self.assertGreater(labels("cat")[0][1], 0.5)
|
||||
|
||||
def testRTD(self):
|
||||
"""
|
||||
Test training a language model with replaced token detection
|
||||
"""
|
||||
|
||||
# Save directory
|
||||
output = os.path.join(tempfile.gettempdir(), "trainer.rtd")
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, _ = trainer("hf-internal-testing/tiny-random-electra", self.data, task="token-detection", output_dir=output)
|
||||
|
||||
# Test model completed successfully
|
||||
self.assertIsNotNone(model)
|
||||
|
||||
# Test output directories exist
|
||||
self.assertTrue(os.path.exists(os.path.join(output, "generator")))
|
||||
self.assertTrue(os.path.exists(os.path.join(output, "discriminator")))
|
||||
|
||||
def testSeqSeq(self):
|
||||
"""
|
||||
Test training a sequence-sequence model
|
||||
"""
|
||||
|
||||
data = [
|
||||
{"source": "Running again", "target": "Sleeping again"},
|
||||
{"source": "Run", "target": "Sleep"},
|
||||
{"source": "running", "target": "sleeping"},
|
||||
]
|
||||
|
||||
trainer = HFTrainer()
|
||||
model, tokenizer = trainer("t5-small", data, task="sequence-sequence", prefix="translate Run to Sleep: ", learning_rate=1e-3)
|
||||
|
||||
# Run run-sleep translation
|
||||
sequences = Sequences((model, tokenizer))
|
||||
result = sequences("translate Run to Sleep: run")
|
||||
self.assertEqual(result.lower(), "sleep")
|
||||
@@ -0,0 +1,503 @@
|
||||
"""
|
||||
Keyword scoring tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from txtai.scoring import Normalize, ScoringFactory, Scoring
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestKeyword(unittest.TestCase):
|
||||
"""
|
||||
Sparse keyword scoring tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"wins wins wins",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
cls.data = [(uid, x, None) for uid, x in enumerate(cls.data)]
|
||||
|
||||
def testBM25(self):
|
||||
"""
|
||||
Test bm25
|
||||
"""
|
||||
|
||||
self.runTests("bm25")
|
||||
|
||||
def testCustom(self):
|
||||
"""
|
||||
Test custom method
|
||||
"""
|
||||
|
||||
self.runTests("txtai.scoring.BM25")
|
||||
|
||||
def testCustomNotFound(self):
|
||||
"""
|
||||
Test unresolvable custom method
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
ScoringFactory.create("notfound.scoring")
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
scoring = Scoring()
|
||||
|
||||
self.assertRaises(NotImplementedError, scoring.insert, None, None)
|
||||
self.assertRaises(NotImplementedError, scoring.delete, None)
|
||||
self.assertRaises(NotImplementedError, scoring.weights, None)
|
||||
self.assertRaises(NotImplementedError, scoring.search, None, None)
|
||||
self.assertRaises(NotImplementedError, scoring.batchsearch, None, None, None)
|
||||
self.assertRaises(NotImplementedError, scoring.count)
|
||||
self.assertRaises(NotImplementedError, scoring.load, None)
|
||||
self.assertRaises(NotImplementedError, scoring.save, None)
|
||||
self.assertRaises(NotImplementedError, scoring.close)
|
||||
self.assertRaises(NotImplementedError, scoring.issparse)
|
||||
self.assertRaises(NotImplementedError, scoring.isnormalized)
|
||||
self.assertRaises(NotImplementedError, scoring.isbayes)
|
||||
|
||||
@patch("sqlalchemy.orm.Query.params")
|
||||
def testPGText(self, query):
|
||||
"""
|
||||
Test PGText
|
||||
"""
|
||||
|
||||
# Mock database query
|
||||
query.return_value = [(3, 1.0)]
|
||||
|
||||
# Create scoring
|
||||
path = os.path.join(tempfile.gettempdir(), "pgtext.sqlite")
|
||||
scoring = ScoringFactory.create({"method": "pgtext", "url": f"sqlite:///{path}", "schema": "txtai"})
|
||||
scoring.index((uid, {"text": text}, tags) for uid, text, tags in self.data)
|
||||
|
||||
# Run search and validate correct result returned
|
||||
index, _ = scoring.search("bear", 1)[0]
|
||||
self.assertEqual(index, 3)
|
||||
|
||||
# Run batch search
|
||||
index, _ = scoring.batchsearch(["bear"], 1)[0][0]
|
||||
self.assertEqual(index, 3)
|
||||
|
||||
# Validate save/load/delete
|
||||
scoring.save(None)
|
||||
scoring.load(None)
|
||||
|
||||
# Validate count
|
||||
self.assertEqual(scoring.count(), len(self.data))
|
||||
|
||||
# Test delete
|
||||
scoring.delete([0])
|
||||
self.assertEqual(scoring.count(), len(self.data) - 1)
|
||||
|
||||
# PGText is a normalized sparse index
|
||||
self.assertTrue(scoring.issparse() and scoring.isnormalized() and not scoring.isbayes())
|
||||
self.assertIsNone(scoring.weights("This is a test".split()))
|
||||
|
||||
# Close scoring
|
||||
scoring.close()
|
||||
|
||||
def testSIF(self):
|
||||
"""
|
||||
Test sif
|
||||
"""
|
||||
|
||||
self.runTests("sif")
|
||||
|
||||
def testTFIDF(self):
|
||||
"""
|
||||
Test tfidf
|
||||
"""
|
||||
|
||||
self.runTests("tfidf")
|
||||
|
||||
def runTests(self, method):
|
||||
"""
|
||||
Runs a series of tests for a scoring method.
|
||||
|
||||
Args:
|
||||
method: scoring method
|
||||
"""
|
||||
|
||||
config = {"method": method}
|
||||
|
||||
self.index(config)
|
||||
self.upsert(config)
|
||||
self.weights(config)
|
||||
self.search(config)
|
||||
self.delete(config)
|
||||
self.normalize(config)
|
||||
self.content(config)
|
||||
self.empty(config)
|
||||
self.copy(config)
|
||||
self.settings(config)
|
||||
self.tokenization(config)
|
||||
|
||||
def index(self, config, data=None):
|
||||
"""
|
||||
Test scoring index method.
|
||||
|
||||
Args:
|
||||
config: scoring config
|
||||
data: data to index with scoring method
|
||||
|
||||
Returns:
|
||||
scoring
|
||||
"""
|
||||
|
||||
# Derive input data
|
||||
data = data if data else self.data
|
||||
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index(data)
|
||||
|
||||
keys = [k for k, v in sorted(scoring.idf.items(), key=lambda x: x[1])]
|
||||
|
||||
# Test count
|
||||
self.assertEqual(scoring.count(), len(data))
|
||||
|
||||
# Win should be lowest score
|
||||
self.assertEqual(keys[0], "wins")
|
||||
|
||||
# Test save/load
|
||||
self.assertIsNotNone(self.save(scoring, config, f"scoring.{config['method']}.index"))
|
||||
|
||||
# Test search returns none when terms disabled (default)
|
||||
self.assertIsNone(scoring.search("query"))
|
||||
|
||||
return scoring
|
||||
|
||||
def upsert(self, config):
|
||||
"""
|
||||
Test scoring upsert method
|
||||
"""
|
||||
|
||||
scoring = ScoringFactory.create({**config, **{"tokenizer": {"alphanum": True, "stopwords": True}}})
|
||||
scoring.upsert(self.data)
|
||||
|
||||
# Test count
|
||||
self.assertEqual(scoring.count(), len(self.data))
|
||||
|
||||
# Test stop word is removed
|
||||
self.assertFalse("and" in scoring.idf)
|
||||
|
||||
def save(self, scoring, config, name):
|
||||
"""
|
||||
Test scoring index save/load.
|
||||
|
||||
Args:
|
||||
scoring: scoring index
|
||||
config: scoring config
|
||||
name: output file name
|
||||
|
||||
Returns:
|
||||
scoring
|
||||
"""
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "scoring")
|
||||
os.makedirs(index, exist_ok=True)
|
||||
|
||||
# Save scoring instance
|
||||
scoring.save(f"{index}/{name}")
|
||||
|
||||
# Reload scoring instance
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.load(f"{index}/{name}")
|
||||
|
||||
return scoring
|
||||
|
||||
def weights(self, config):
|
||||
"""
|
||||
Test standard and tag weighted scores.
|
||||
|
||||
Args:
|
||||
config: scoring config
|
||||
"""
|
||||
|
||||
document = (1, ["bear", "wins"], None)
|
||||
|
||||
scoring = self.index(config)
|
||||
weights = scoring.weights(document[1])
|
||||
|
||||
# Default weights
|
||||
self.assertNotEqual(weights[0], weights[1])
|
||||
|
||||
data = self.data[:]
|
||||
|
||||
uid, text, _ = data[3]
|
||||
data[3] = (uid, text, "wins")
|
||||
|
||||
scoring = self.index(config, data)
|
||||
weights = scoring.weights(document[1])
|
||||
|
||||
# Modified weights
|
||||
self.assertEqual(weights[0], weights[1])
|
||||
|
||||
def search(self, config):
|
||||
"""
|
||||
Test scoring search.
|
||||
|
||||
Args:
|
||||
config: scoring config
|
||||
"""
|
||||
|
||||
# Create combined config
|
||||
config = {**config, **{"terms": True}}
|
||||
|
||||
# Create scoring instance
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index(self.data)
|
||||
|
||||
# Run search and validate correct result returned
|
||||
index, _ = scoring.search("bear", 1)[0]
|
||||
self.assertEqual(index, 3)
|
||||
|
||||
# Run batch search
|
||||
index, _ = scoring.batchsearch(["bear"], 1)[0][0]
|
||||
self.assertEqual(index, 3)
|
||||
|
||||
# Run wildcard search
|
||||
index, _ = scoring.search("bea*", 1)[0]
|
||||
self.assertEqual(index, 3)
|
||||
|
||||
# Test save/reload
|
||||
self.save(scoring, config, f"scoring.{config['method']}.search")
|
||||
|
||||
# Re-run search and validate correct result returned
|
||||
index, _ = scoring.search("bear", 1)[0]
|
||||
self.assertEqual(index, 3)
|
||||
|
||||
def delete(self, config):
|
||||
"""
|
||||
Test delete.
|
||||
"""
|
||||
|
||||
# Create combined config
|
||||
config = {**config, **{"terms": True, "content": True}}
|
||||
|
||||
# Create scoring instance
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index(self.data)
|
||||
|
||||
# Run search and validate correct result returned
|
||||
index = scoring.search("bear", 1)[0]["id"]
|
||||
|
||||
# Delete result and validate the query no longer returns results
|
||||
scoring.delete([index])
|
||||
self.assertFalse(scoring.search("bear", 1))
|
||||
|
||||
# Save and validate count
|
||||
self.save(scoring, config, f"scoring.{config['method']}.delete")
|
||||
self.assertEqual(scoring.count(), len(self.data) - 1)
|
||||
|
||||
def normalize(self, config):
|
||||
"""
|
||||
Test scoring search with normalized scores.
|
||||
|
||||
Args:
|
||||
method: scoring method
|
||||
"""
|
||||
|
||||
# Default normalization
|
||||
scoring = ScoringFactory.create({**config, **{"terms": True, "normalize": True}})
|
||||
scoring.index(self.data)
|
||||
|
||||
# Run search and validate correct result returned
|
||||
index, score = scoring.search(self.data[3][1], 1)[0]
|
||||
self.assertEqual(index, 3)
|
||||
self.assertEqual(score, 1.0)
|
||||
|
||||
# Bayesian normalization with default dynamic alpha/beta settings
|
||||
baseline = ScoringFactory.create({**config, **{"terms": True}})
|
||||
baseline.index(self.data)
|
||||
|
||||
scoring = ScoringFactory.create({**config, **{"terms": True, "normalize": "bayes"}})
|
||||
scoring.index(self.data)
|
||||
|
||||
query = "wins"
|
||||
base = baseline.search(query, 3)
|
||||
bayes = scoring.search(query, 3)
|
||||
|
||||
# Bayesian normalization should preserve ranking order while mapping scores to [0, 1]
|
||||
self.assertEqual([uid for uid, _ in base], [uid for uid, _ in bayes])
|
||||
self.assertTrue(all(0.0 <= score <= 1.0 for _, score in bayes))
|
||||
|
||||
# BB25 alias should resolve to Bayesian normalization
|
||||
scoring = ScoringFactory.create({**config, **{"terms": True, "normalize": "bb25"}})
|
||||
scoring.index(self.data)
|
||||
bb25 = scoring.search(query, 3)
|
||||
self.assertEqual([uid for uid, _ in base], [uid for uid, _ in bb25])
|
||||
self.assertTrue(all(0.0 <= score <= 1.0 for _, score in bb25))
|
||||
|
||||
# BB25 candidate-set behavior: zero scores remain 0, positive scores are transformed
|
||||
normalizer = Normalize("bb25")
|
||||
scores = normalizer([(0, 0.0), (1, 1.0), (2, 2.0)], scoring.avgscore)
|
||||
self.assertEqual(scores[0][1], 0.0)
|
||||
self.assertGreater(scores[1][1], 0.0)
|
||||
self.assertGreater(scores[2][1], scores[1][1])
|
||||
|
||||
# Test negative scores
|
||||
scores = normalizer([(0, -100.0)], scoring.avgscore)
|
||||
self.assertEqual(scores[0][1], 0.0)
|
||||
|
||||
# Bayesian normalization with custom parameters
|
||||
config = {**config, **{"terms": True, "normalize": {"method": "bayes", "alpha": 2.0}}}
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index(self.data)
|
||||
|
||||
custom = scoring.search(query, 3)
|
||||
self.assertEqual([uid for uid, _ in base], [uid for uid, _ in custom])
|
||||
self.assertTrue(all(0.0 <= score <= 1.0 for _, score in custom))
|
||||
|
||||
def content(self, config):
|
||||
"""
|
||||
Test scoring search with content.
|
||||
|
||||
Args:
|
||||
config: scoring config
|
||||
"""
|
||||
|
||||
scoring = ScoringFactory.create({**config, **{"terms": True, "content": True}})
|
||||
scoring.index(self.data)
|
||||
|
||||
# Test text with content
|
||||
text = "Great news today"
|
||||
scoring.index([(scoring.total, text, None)])
|
||||
|
||||
# Run search and validate correct result returned
|
||||
result = scoring.search("great news", 1)[0]["text"]
|
||||
self.assertEqual(result, text)
|
||||
|
||||
# Test reading text from dictionary
|
||||
text = "Feel good story: baby panda born"
|
||||
scoring.index([(scoring.total, {"text": text}, None)])
|
||||
|
||||
# Run search and validate correct result returned
|
||||
result = scoring.search("feel good story", 1)[0]["text"]
|
||||
self.assertEqual(result, text)
|
||||
|
||||
def empty(self, config):
|
||||
"""
|
||||
Test scoring index properly handles an index call when no data present.
|
||||
|
||||
Args:
|
||||
config: scoring config
|
||||
"""
|
||||
|
||||
# Create scoring index with no data
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index([])
|
||||
|
||||
# Assert index call returns and index has a count of 0
|
||||
self.assertEqual(scoring.total, 0)
|
||||
|
||||
def copy(self, config):
|
||||
"""
|
||||
Test scoring index copy method.
|
||||
"""
|
||||
|
||||
# Create scoring instance
|
||||
scoring = ScoringFactory.create({**config, **{"terms": True}})
|
||||
scoring.index(self.data)
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "scoring")
|
||||
os.makedirs(index, exist_ok=True)
|
||||
|
||||
# Create file to test replacing existing file
|
||||
path = f"{index}/scoring.{config['method']}.copy"
|
||||
with open(f"{index}.terms", "w", encoding="utf-8") as f:
|
||||
f.write("TEST")
|
||||
|
||||
# Save scoring instance
|
||||
scoring.save(path)
|
||||
self.assertTrue(os.path.exists(path))
|
||||
|
||||
@patch("sys.byteorder", "big")
|
||||
def settings(self, config):
|
||||
"""
|
||||
Test various settings.
|
||||
|
||||
Args:
|
||||
config: scoring config
|
||||
"""
|
||||
|
||||
# Create combined config
|
||||
config = {**config, **{"terms": {"cachelimit": 0, "cutoff": 0.25, "wal": True}}}
|
||||
|
||||
# Create scoring instance
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index(self.data)
|
||||
|
||||
# Save/load index
|
||||
self.save(scoring, config, f"scoring.{config['method']}.settings")
|
||||
|
||||
index, _ = scoring.search("bear bear bear wins", 1)[0]
|
||||
self.assertEqual(index, 3)
|
||||
|
||||
# Save to same path
|
||||
self.save(scoring, config, f"scoring.{config['method']}.settings")
|
||||
|
||||
# Save to different path
|
||||
self.save(scoring, config, f"scoring.{config['method']}.move")
|
||||
|
||||
# Validate counts
|
||||
self.assertEqual(scoring.count(), len(self.data))
|
||||
|
||||
def tokenization(self, config):
|
||||
"""
|
||||
Test tokenization methods.
|
||||
|
||||
Args:
|
||||
config: scoring config
|
||||
"""
|
||||
|
||||
# Test whitespace tokenization
|
||||
config = {**config, **{"terms": True, "tokenizer": {"whitespace": True}}}
|
||||
|
||||
# Create scoring instance
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index([(0, "abc-def-123", None)])
|
||||
|
||||
self.assertEqual(scoring.search("abc-def-123")[0][0], 0)
|
||||
|
||||
# Test regular expression tokenization
|
||||
config = {**config, **{"tokenizer": {"regexp": r"\w{5,}"}}}
|
||||
|
||||
# Create scoring instance
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index([(0, "hello test", None)])
|
||||
|
||||
self.assertEqual(scoring.search("hello")[0][0], 0)
|
||||
self.assertFalse(scoring.search("test"))
|
||||
|
||||
# Test ngram tokenization
|
||||
ngrams = {"ngrams": 3, "lpad": " ", "rpad": " ", "unique": True}
|
||||
config = {**config, **{"tokenizer": {"ngrams": ngrams}}}
|
||||
|
||||
# Create scoring instance
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index([(0, "hello test", None)])
|
||||
|
||||
self.assertEqual(scoring.search("hello")[0][0], 0)
|
||||
@@ -0,0 +1,208 @@
|
||||
"""
|
||||
Sparse module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import platform
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from txtai.scoring import ScoringFactory
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestSparse(unittest.TestCase):
|
||||
"""
|
||||
Sparse vector scoring tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
cls.data = [
|
||||
"US tops 5 million confirmed virus cases",
|
||||
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
||||
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
||||
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
||||
"Maine man wins $1M from $25 lottery ticket",
|
||||
"Make huge profits without work, earn up to $100,000 a day",
|
||||
]
|
||||
|
||||
cls.data = [(uid, x, None) for uid, x in enumerate(cls.data)]
|
||||
|
||||
def testGeneral(self):
|
||||
"""
|
||||
Test general sparse vector operations
|
||||
"""
|
||||
|
||||
# Models cache
|
||||
models = {}
|
||||
|
||||
# Test sparse scoring
|
||||
scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}, models=models)
|
||||
scoring.index((uid, {"text": text}, tags) for uid, text, tags in self.data)
|
||||
|
||||
# Run search and validate correct result returned
|
||||
index, _ = scoring.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(index, 4)
|
||||
|
||||
# Run batch search
|
||||
index, _ = scoring.batchsearch(["lottery ticket"], 1)[0][0]
|
||||
self.assertEqual(index, 4)
|
||||
|
||||
# Validate count
|
||||
self.assertEqual(scoring.count(), len(self.data))
|
||||
|
||||
# Test delete
|
||||
scoring.delete([4])
|
||||
self.assertEqual(scoring.count(), len(self.data) - 1)
|
||||
|
||||
# Run search after delete
|
||||
index, _ = scoring.search("lottery ticket", 1)[0]
|
||||
self.assertEqual(index, 5)
|
||||
|
||||
# Sparse vectors is a normalized sparse index
|
||||
self.assertTrue(scoring.issparse() and scoring.isnormalized() and not scoring.isbayes())
|
||||
self.assertIsNone(scoring.weights("This is a test".split()))
|
||||
|
||||
# Close scoring
|
||||
scoring.close()
|
||||
|
||||
# Test model caching
|
||||
scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq"}, models=models)
|
||||
self.assertIsNotNone(scoring.model)
|
||||
scoring.close()
|
||||
|
||||
def testEmpty(self):
|
||||
"""
|
||||
Test empty sparse vectors
|
||||
"""
|
||||
|
||||
scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq"})
|
||||
scoring.upsert((uid, {"text": text}, tags) for uid, text, tags in self.data)
|
||||
self.assertEqual(scoring.count(), len(self.data))
|
||||
|
||||
@unittest.skipIf(platform.system() == "Darwin", "Torch memory sharing not supported on macOS")
|
||||
@patch("torch.cuda.device_count")
|
||||
def testGPU(self, count):
|
||||
"""
|
||||
Test sparse vectors with GPU encoding
|
||||
"""
|
||||
|
||||
# Mock accelerator count
|
||||
count.return_value = 2
|
||||
|
||||
# Test multiple gpus
|
||||
scoring = ScoringFactory.create({"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq", "gpu": "all"})
|
||||
self.assertIsNotNone(scoring)
|
||||
scoring.close()
|
||||
|
||||
def testBayes(self):
|
||||
"""
|
||||
Test BB25 Bayesian normalization for sparse scoring
|
||||
"""
|
||||
|
||||
config = {
|
||||
"method": "sparse",
|
||||
"path": "sparse-encoder-testing/splade-bert-tiny-nq",
|
||||
"normalize": "bb25",
|
||||
}
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index((uid, {"text": text}, tags) for uid, text, tags in self.data)
|
||||
|
||||
# Verify Bayesian mode flags
|
||||
self.assertTrue(scoring.isbayes())
|
||||
self.assertTrue(scoring.isnormalized())
|
||||
|
||||
# Search and validate scores are calibrated probabilities in [0, 1]
|
||||
results = scoring.search("lottery ticket", 3)
|
||||
self.assertGreater(len(results), 0)
|
||||
for _, score in results:
|
||||
self.assertGreaterEqual(score, 0.0)
|
||||
self.assertLessEqual(score, 1.0)
|
||||
|
||||
# Batch search
|
||||
results = scoring.batchsearch(["lottery ticket", "ice shelf"], 3)
|
||||
self.assertEqual(len(results), 2)
|
||||
for query_results in results:
|
||||
for _, score in query_results:
|
||||
self.assertGreaterEqual(score, 0.0)
|
||||
self.assertLessEqual(score, 1.0)
|
||||
|
||||
scoring.close()
|
||||
|
||||
def testBayesDict(self):
|
||||
"""
|
||||
Test BB25 normalization with dict config
|
||||
"""
|
||||
|
||||
config = {
|
||||
"method": "sparse",
|
||||
"path": "sparse-encoder-testing/splade-bert-tiny-nq",
|
||||
"normalize": {"method": "bb25", "alpha": 2.0},
|
||||
}
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index((uid, {"text": text}, tags) for uid, text, tags in self.data)
|
||||
|
||||
self.assertTrue(scoring.isbayes())
|
||||
|
||||
results = scoring.search("lottery ticket", 3)
|
||||
self.assertGreater(len(results), 0)
|
||||
for _, score in results:
|
||||
self.assertGreaterEqual(score, 0.0)
|
||||
self.assertLessEqual(score, 1.0)
|
||||
|
||||
scoring.close()
|
||||
|
||||
def testBayesNonBayes(self):
|
||||
"""
|
||||
Test that non-Bayesian string normalize values do not activate Bayesian mode
|
||||
"""
|
||||
|
||||
config = {
|
||||
"method": "sparse",
|
||||
"path": "sparse-encoder-testing/splade-bert-tiny-nq",
|
||||
"normalize": "default",
|
||||
}
|
||||
scoring = ScoringFactory.create(config)
|
||||
self.assertFalse(scoring.isbayes())
|
||||
scoring.close()
|
||||
|
||||
def testIVFFlat(self):
|
||||
"""
|
||||
Test sparse vectors with IVFFlat clustering
|
||||
"""
|
||||
|
||||
# Expand dataset
|
||||
data = self.data * 1000
|
||||
|
||||
# Test higher volume IVFFlat index with clustering
|
||||
config = {
|
||||
"method": "sparse",
|
||||
"vectormethod": "sentence-transformers",
|
||||
"path": "sparse-encoder-testing/splade-bert-tiny-nq",
|
||||
"ivfsparse": {"sample": 1.0},
|
||||
}
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.index((uid, {"text": text}, tags) for uid, text, tags in data)
|
||||
|
||||
# Generate temp file path
|
||||
index = os.path.join(tempfile.gettempdir(), "scoring")
|
||||
os.makedirs(index, exist_ok=True)
|
||||
|
||||
# Save scoring instance
|
||||
scoring.save(f"{index}/scoring.sparse.index")
|
||||
|
||||
# Reload scoring instance
|
||||
scoring = ScoringFactory.create(config)
|
||||
scoring.load(f"{index}/scoring.sparse.index")
|
||||
|
||||
# Run search and validate correct result returned
|
||||
results = scoring.search("lottery ticket", 1)
|
||||
self.assertGreater(len(results), 0)
|
||||
scoring.close()
|
||||
@@ -0,0 +1,61 @@
|
||||
"""
|
||||
Serialize module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from txtai.serialize import Serialize, SerializeFactory
|
||||
|
||||
|
||||
class TestSerialize(unittest.TestCase):
|
||||
"""
|
||||
Serialize tests.
|
||||
"""
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
serialize = Serialize()
|
||||
|
||||
self.assertRaises(NotImplementedError, serialize.loadstream, None)
|
||||
self.assertRaises(NotImplementedError, serialize.savestream, None, None)
|
||||
self.assertRaises(NotImplementedError, serialize.loadbytes, None)
|
||||
self.assertRaises(NotImplementedError, serialize.savebytes, None)
|
||||
|
||||
def testMessagePack(self):
|
||||
"""
|
||||
Test MessagePack encoder
|
||||
"""
|
||||
|
||||
serializer = SerializeFactory.create()
|
||||
self.assertEqual(serializer.loadbytes(serializer.savebytes("test")), "test")
|
||||
|
||||
def testPickleDisabled(self):
|
||||
"""
|
||||
Test disabled pickle serialization
|
||||
"""
|
||||
|
||||
# Validate an error is raised
|
||||
with self.assertRaises(ValueError):
|
||||
serializer = SerializeFactory.create("pickle", allowpickle=True)
|
||||
data = serializer.savebytes("Test")
|
||||
|
||||
serializer = SerializeFactory.create("pickle")
|
||||
serializer.loadbytes(data)
|
||||
|
||||
@patch.dict(os.environ, {"ALLOW_PICKLE": "True"})
|
||||
def testPickleEnabled(self):
|
||||
"""
|
||||
Test enabled pickle serialization
|
||||
"""
|
||||
|
||||
# Validate a warning is raised
|
||||
with self.assertWarns(RuntimeWarning):
|
||||
serializer = SerializeFactory.create("pickle")
|
||||
data = serializer.savebytes("Test")
|
||||
serializer.loadbytes(data)
|
||||
@@ -0,0 +1,51 @@
|
||||
"""
|
||||
Custom module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.vectors import VectorsFactory
|
||||
|
||||
|
||||
class TestCustom(unittest.TestCase):
|
||||
"""
|
||||
Custom vectors tests
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create custom vectors instance.
|
||||
"""
|
||||
|
||||
cls.model = VectorsFactory.create({"method": "txtai.vectors.HFVectors", "path": "sentence-transformers/nli-mpnet-base-v2"}, None)
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test transformers indexing
|
||||
"""
|
||||
|
||||
# Generate enough volume to test batching
|
||||
documents = [(x, "This is a test", None) for x in range(1000)]
|
||||
|
||||
ids, dimension, batches, stream = self.model.index(documents)
|
||||
|
||||
self.assertEqual(len(ids), 1000)
|
||||
self.assertEqual(dimension, 768)
|
||||
self.assertEqual(batches, 2)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (500, 768))
|
||||
|
||||
def testNotFound(self):
|
||||
"""
|
||||
Test unresolvable vector backend
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
VectorsFactory.create({"method": "notfound.vectors"})
|
||||
@@ -0,0 +1,54 @@
|
||||
"""
|
||||
External module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.vectors import External, VectorsFactory
|
||||
|
||||
|
||||
class TestExternal(unittest.TestCase):
|
||||
"""
|
||||
External vectors tests
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create External vectors instance.
|
||||
"""
|
||||
|
||||
cls.model = VectorsFactory.create({"method": "external"}, None)
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test indexing with external vectors
|
||||
"""
|
||||
|
||||
# Generate dummy data
|
||||
data = np.random.rand(1000, 768).astype(np.float32)
|
||||
|
||||
# Generate enough volume to test batching
|
||||
documents = [(x, data[x], None) for x in range(1000)]
|
||||
|
||||
ids, dimension, batches, stream = self.model.index(documents)
|
||||
|
||||
self.assertEqual(len(ids), 1000)
|
||||
self.assertEqual(dimension, 768)
|
||||
self.assertEqual(batches, 2)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (500, 768))
|
||||
|
||||
def testMethod(self):
|
||||
"""
|
||||
Test method is derived when transform function passed
|
||||
"""
|
||||
|
||||
model = VectorsFactory.create({"transform": lambda x: x}, None)
|
||||
self.assertTrue(isinstance(model, External))
|
||||
@@ -0,0 +1,99 @@
|
||||
"""
|
||||
Huggingface module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.vectors import VectorsFactory
|
||||
|
||||
|
||||
class TestHFVectors(unittest.TestCase):
|
||||
"""
|
||||
HFVectors tests
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create HFVectors instance.
|
||||
"""
|
||||
|
||||
cls.model = VectorsFactory.create({"path": "sentence-transformers/nli-mpnet-base-v2"}, None)
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test transformers indexing
|
||||
"""
|
||||
|
||||
# Generate enough volume to test batching
|
||||
documents = [(x, "This is a test", None) for x in range(1000)]
|
||||
|
||||
ids, dimension, batches, stream = self.model.index(documents)
|
||||
|
||||
self.assertEqual(len(ids), 1000)
|
||||
self.assertEqual(dimension, 768)
|
||||
self.assertEqual(batches, 2)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (500, 768))
|
||||
|
||||
def testText(self):
|
||||
"""
|
||||
Test transformers text conversion
|
||||
"""
|
||||
|
||||
self.model.tokenize = True
|
||||
self.assertEqual(self.model.prepare("Y 123 This is a test!"), "test")
|
||||
self.assertEqual(self.model.prepare(["This", "is", "a", "test"]), "This is a test")
|
||||
|
||||
self.model.tokenize = False
|
||||
self.assertEqual(self.model.prepare("Y 123 This is a test!"), "Y 123 This is a test!")
|
||||
self.assertEqual(self.model.prepare(["This", "is", "a", "test"]), "This is a test")
|
||||
|
||||
def testTransform(self):
|
||||
"""
|
||||
Test transformers transform
|
||||
"""
|
||||
|
||||
# Sample documents: one where tokenizer changes text and one with no changes to text
|
||||
documents = [(0, "This is a test and has no tokenization", None), (1, "test tokenization", None)]
|
||||
|
||||
# Run with tokenization enabled
|
||||
self.model.tokenize = True
|
||||
embeddings1 = [self.model.transform(d) for d in documents]
|
||||
|
||||
# Run with tokenization disabled
|
||||
self.model.tokenize = False
|
||||
embeddings2 = [self.model.transform(d) for d in documents]
|
||||
|
||||
self.assertFalse(np.array_equal(embeddings1[0], embeddings2[0]))
|
||||
self.assertTrue(np.array_equal(embeddings1[1], embeddings2[1]))
|
||||
|
||||
def testTransformArray(self):
|
||||
"""
|
||||
Test transformers skips transforming NumPy arrays
|
||||
"""
|
||||
|
||||
# Generate data and run through vector model
|
||||
data1 = np.random.rand(5, 5).astype(np.float32)
|
||||
data2 = self.model.transform((0, data1, None))
|
||||
|
||||
# Test transform method returns original data
|
||||
self.assertTrue(np.array_equal(data1, data2))
|
||||
|
||||
def testTransformLong(self):
|
||||
"""
|
||||
Test transformers transform on long text
|
||||
"""
|
||||
|
||||
# Sample documents: short text and longer text
|
||||
documents = [(0, "This is long text " * 512, None), (1, "This is short text", None)]
|
||||
|
||||
# Run transform and ensure it completes without errors
|
||||
embeddings = [self.model.transform(d) for d in documents]
|
||||
self.assertIsNotNone(embeddings)
|
||||
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
LiteLLM module tests
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from threading import Thread
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.vectors import VectorsFactory
|
||||
|
||||
|
||||
class RequestHandler(BaseHTTPRequestHandler):
|
||||
"""
|
||||
Test HTTP handler.
|
||||
"""
|
||||
|
||||
def do_POST(self):
|
||||
"""
|
||||
POST request handler.
|
||||
"""
|
||||
|
||||
# Generate mock response
|
||||
response = [[0.0] * 768]
|
||||
response = json.dumps(response).encode("utf-8")
|
||||
|
||||
self.send_response(200)
|
||||
self.send_header("content-type", "application/json")
|
||||
self.send_header("content-length", len(response))
|
||||
self.end_headers()
|
||||
|
||||
self.wfile.write(response)
|
||||
self.wfile.flush()
|
||||
|
||||
|
||||
class TestLiteLLM(unittest.TestCase):
|
||||
"""
|
||||
LiteLLM vectors tests
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create mock http server.
|
||||
"""
|
||||
|
||||
cls.httpd = HTTPServer(("127.0.0.1", 8004), RequestHandler)
|
||||
|
||||
server = Thread(target=cls.httpd.serve_forever, daemon=True)
|
||||
server.start()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""
|
||||
Shutdown mock http server.
|
||||
"""
|
||||
|
||||
cls.httpd.shutdown()
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test indexing with LiteLLM vectors
|
||||
"""
|
||||
|
||||
# LiteLLM vectors instance
|
||||
model = VectorsFactory.create(
|
||||
{"path": "huggingface/sentence-transformers/all-MiniLM-L6-v2", "vectors": {"api_base": "http://127.0.0.1:8004"}}, None
|
||||
)
|
||||
|
||||
ids, dimension, batches, stream = model.index([(0, "test", None)])
|
||||
|
||||
self.assertEqual(len(ids), 1)
|
||||
self.assertEqual(dimension, 768)
|
||||
self.assertEqual(batches, 1)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (1, 768))
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
LiteRT module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.vectors import VectorsFactory
|
||||
|
||||
|
||||
class TestLiteRT(unittest.TestCase):
|
||||
"""
|
||||
LiteRT vectors tests
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create LiteRT instance.
|
||||
"""
|
||||
|
||||
cls.model = VectorsFactory.create(
|
||||
{"path": "neuml/bert-hash-nano-embeddings-litert/bert-hash-nano-embeddings-int4.tflite", "gpu": False}, None
|
||||
)
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test indexing with LiteRT vectors
|
||||
"""
|
||||
|
||||
ids, dimension, batches, stream = self.model.index([(0, "test", None)])
|
||||
|
||||
self.assertEqual(len(ids), 1)
|
||||
self.assertEqual(dimension, 128)
|
||||
self.assertEqual(batches, 1)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (1, 128))
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Llama module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.vectors import VectorsFactory
|
||||
|
||||
|
||||
class TestLlamaCpp(unittest.TestCase):
|
||||
"""
|
||||
llama.cpp vectors tests
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create LlamaCpp instance.
|
||||
"""
|
||||
|
||||
cls.model = VectorsFactory.create({"path": "nomic-ai/nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.Q2_K.gguf"}, None)
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test indexing with LlamaCpp vectors
|
||||
"""
|
||||
|
||||
ids, dimension, batches, stream = self.model.index([(0, "test", None)])
|
||||
|
||||
self.assertEqual(len(ids), 1)
|
||||
self.assertEqual(dimension, 768)
|
||||
self.assertEqual(batches, 1)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (1, 768))
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Model2Vec module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.vectors import VectorsFactory
|
||||
|
||||
|
||||
class TestModel2Vec(unittest.TestCase):
|
||||
"""
|
||||
Model2vec vectors tests
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Create Model2Vec instance.
|
||||
"""
|
||||
|
||||
cls.model = VectorsFactory.create({"path": "minishlab/potion-base-8M"}, None)
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test indexing with Model2Vec vectors
|
||||
"""
|
||||
|
||||
ids, dimension, batches, stream = self.model.index([(0, "test", None)])
|
||||
|
||||
self.assertEqual(len(ids), 1)
|
||||
self.assertEqual(dimension, 256)
|
||||
self.assertEqual(batches, 1)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (1, 256))
|
||||
@@ -0,0 +1,61 @@
|
||||
"""
|
||||
Sentence Transformers module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import platform
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.vectors import VectorsFactory
|
||||
|
||||
|
||||
class TestSTVectors(unittest.TestCase):
|
||||
"""
|
||||
STVectors tests
|
||||
"""
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test indexing with sentence-transformers vectors
|
||||
"""
|
||||
|
||||
model = VectorsFactory.create({"method": "sentence-transformers", "path": "paraphrase-MiniLM-L3-v2"}, None)
|
||||
ids, dimension, batches, stream = model.index([(0, "test", None)])
|
||||
|
||||
self.assertEqual(len(ids), 1)
|
||||
self.assertEqual(dimension, 384)
|
||||
self.assertEqual(batches, 1)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (1, 384))
|
||||
|
||||
@unittest.skipIf(platform.system() == "Darwin", "Torch memory sharing not supported on macOS")
|
||||
@patch("torch.cuda.device_count")
|
||||
def testMultiGPU(self, count):
|
||||
"""
|
||||
Test multiple gpu encoding
|
||||
"""
|
||||
|
||||
# Mock accelerator count
|
||||
count.return_value = 2
|
||||
|
||||
model = VectorsFactory.create({"method": "sentence-transformers", "path": "paraphrase-MiniLM-L3-v2", "gpu": "all"}, None)
|
||||
ids, dimension, batches, stream = model.index([(0, "test", None)])
|
||||
|
||||
self.assertEqual(len(ids), 1)
|
||||
self.assertEqual(dimension, 384)
|
||||
self.assertEqual(batches, 1)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (1, 384))
|
||||
|
||||
# Close the multiprocessing pool
|
||||
model.close()
|
||||
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
Vectors module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from txtai.vectors import Vectors, Recovery
|
||||
|
||||
|
||||
class TestVectors(unittest.TestCase):
|
||||
"""
|
||||
Vectors tests.
|
||||
"""
|
||||
|
||||
def testNotImplemented(self):
|
||||
"""
|
||||
Test exceptions for non-implemented methods
|
||||
"""
|
||||
|
||||
vectors = Vectors(None, None, None)
|
||||
|
||||
self.assertRaises(NotImplementedError, vectors.load, None)
|
||||
self.assertRaises(NotImplementedError, vectors.encode, None)
|
||||
|
||||
def testNormalize(self):
|
||||
"""
|
||||
Test batch normalize and single input normalize are equal
|
||||
"""
|
||||
|
||||
vectors = Vectors(None, None, None)
|
||||
|
||||
# Generate data
|
||||
data1 = np.random.rand(5, 5).astype(np.float32)
|
||||
data2 = data1.copy()
|
||||
|
||||
# Keep original data to ensure it changed
|
||||
original = data1.copy()
|
||||
|
||||
# Normalize data
|
||||
vectors.normalize(data1)
|
||||
for x in data2:
|
||||
vectors.normalize(x)
|
||||
|
||||
# Test both data arrays are the same and changed from original
|
||||
self.assertTrue(np.allclose(data1, data2))
|
||||
self.assertFalse(np.allclose(data1, original))
|
||||
|
||||
def testRecovery(self):
|
||||
"""
|
||||
Test vectors recovery failure
|
||||
"""
|
||||
|
||||
# Checkpoint directory
|
||||
checkpoint = os.path.join(tempfile.gettempdir(), "recovery")
|
||||
os.makedirs(checkpoint, exist_ok=True)
|
||||
|
||||
# Create empty file
|
||||
# pylint: disable=R1732
|
||||
f = open(os.path.join(checkpoint, "id"), "w", encoding="utf-8")
|
||||
f.close()
|
||||
|
||||
# Create the recovery instance with an empty checkpoint file
|
||||
recovery = Recovery(checkpoint, "id", np.load)
|
||||
self.assertIsNone(recovery())
|
||||
@@ -0,0 +1,163 @@
|
||||
"""
|
||||
WordVectors module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
from huggingface_hub.errors import HFValidationError
|
||||
from txtai.vectors import VectorsFactory
|
||||
from txtai.vectors.dense.words import create, transform, WordVectors
|
||||
|
||||
|
||||
class TestWordVectors(unittest.TestCase):
|
||||
"""
|
||||
Vectors tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Sets the pretrained model to use
|
||||
"""
|
||||
|
||||
# Test with pretrained glove quantized vectors
|
||||
cls.path = "neuml/glove-6B-quantized"
|
||||
|
||||
@patch("os.cpu_count")
|
||||
def testIndex(self, cpucount):
|
||||
"""
|
||||
Test word vectors indexing
|
||||
"""
|
||||
|
||||
# Mock CPU count
|
||||
cpucount.return_value = 1
|
||||
|
||||
# Generate data
|
||||
documents = [(x, "This is a test", None) for x in range(1000)]
|
||||
|
||||
model = VectorsFactory.create({"path": self.path, "parallel": True}, None)
|
||||
|
||||
ids, dimension, batches, stream = model.index(documents, 1)
|
||||
|
||||
self.assertEqual(len(ids), 1000)
|
||||
self.assertEqual(dimension, 300)
|
||||
self.assertEqual(batches, 1000)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (1, 300))
|
||||
|
||||
@patch("os.cpu_count")
|
||||
def testIndexBatch(self, cpucount):
|
||||
"""
|
||||
Test word vectors indexing with batch size set
|
||||
"""
|
||||
|
||||
# Mock CPU count
|
||||
cpucount.return_value = 1
|
||||
|
||||
# Generate data
|
||||
documents = [(x, "This is a test", None) for x in range(1000)]
|
||||
|
||||
model = VectorsFactory.create({"path": self.path, "parallel": True}, None)
|
||||
|
||||
ids, dimension, batches, stream = model.index(documents, 512)
|
||||
|
||||
self.assertEqual(len(ids), 1000)
|
||||
self.assertEqual(dimension, 300)
|
||||
self.assertEqual(batches, 2)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (512, 300))
|
||||
self.assertEqual(np.load(queue).shape, (488, 300))
|
||||
|
||||
def testIndexSerial(self):
|
||||
"""
|
||||
Test word vector indexing in single process mode
|
||||
"""
|
||||
|
||||
# Generate data
|
||||
documents = [(x, "This is a test", None) for x in range(1000)]
|
||||
|
||||
model = VectorsFactory.create({"path": self.path, "parallel": False}, None)
|
||||
|
||||
ids, dimension, batches, stream = model.index(documents, 1)
|
||||
|
||||
self.assertEqual(len(ids), 1000)
|
||||
self.assertEqual(dimension, 300)
|
||||
self.assertEqual(batches, 1000)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (1, 300))
|
||||
|
||||
def testIndexSerialBatch(self):
|
||||
"""
|
||||
Test word vector indexing in single process mode with batch size set
|
||||
"""
|
||||
|
||||
# Generate data
|
||||
documents = [(x, "This is a test", None) for x in range(1000)]
|
||||
|
||||
model = VectorsFactory.create({"path": self.path, "parallel": False}, None)
|
||||
|
||||
ids, dimension, batches, stream = model.index(documents, 512)
|
||||
|
||||
self.assertEqual(len(ids), 1000)
|
||||
self.assertEqual(dimension, 300)
|
||||
self.assertEqual(batches, 2)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(np.load(queue).shape, (512, 300))
|
||||
self.assertEqual(np.load(queue).shape, (488, 300))
|
||||
|
||||
def testLookup(self):
|
||||
"""
|
||||
Test word vector lookup
|
||||
"""
|
||||
|
||||
model = VectorsFactory.create({"path": self.path}, None)
|
||||
self.assertEqual(model.lookup(["txtai", "embeddings", "sentence"]).shape, (3, 300))
|
||||
|
||||
def testMultiprocess(self):
|
||||
"""
|
||||
Test multiprocess helper methods
|
||||
"""
|
||||
|
||||
create({"path": self.path}, None)
|
||||
|
||||
uid, vector = transform((0, "test", None))
|
||||
self.assertEqual(uid, 0)
|
||||
self.assertEqual(vector.shape, (300,))
|
||||
|
||||
def testNoExist(self):
|
||||
"""
|
||||
Test loading model that doesn't exist
|
||||
"""
|
||||
|
||||
# Test non-existent local path raises an exception
|
||||
with self.assertRaises((IOError, HFValidationError)):
|
||||
VectorsFactory.create({"method": "words", "path": os.path.join(tempfile.gettempdir(), "noexist")}, None)
|
||||
|
||||
# Test non-existent hub path is handled properly
|
||||
self.assertFalse(WordVectors.ismodel("invalid/user/noexist"))
|
||||
|
||||
def testTransform(self):
|
||||
"""
|
||||
Test word vector transform
|
||||
"""
|
||||
|
||||
model = VectorsFactory.create({"path": self.path}, None)
|
||||
self.assertEqual(len(model.transform((None, ["txtai"], None))), 300)
|
||||
@@ -0,0 +1,32 @@
|
||||
"""
|
||||
Sparse Sentence Transformers module tests
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
from txtai.vectors import SparseVectorsFactory
|
||||
from txtai.util import SparseArray
|
||||
|
||||
|
||||
class TestSparseSTVectors(unittest.TestCase):
|
||||
"""
|
||||
SparseSTVectors tests
|
||||
"""
|
||||
|
||||
def testIndex(self):
|
||||
"""
|
||||
Test indexing with sentence-transformers vectors
|
||||
"""
|
||||
|
||||
model = SparseVectorsFactory.create({"method": "sentence-transformers", "path": "sparse-encoder-testing/splade-bert-tiny-nq"})
|
||||
ids, dimension, batches, stream = model.index([(0, "test", None)])
|
||||
|
||||
self.assertEqual(len(ids), 1)
|
||||
self.assertEqual(dimension, 30522)
|
||||
self.assertEqual(batches, 1)
|
||||
self.assertIsNotNone(os.path.exists(stream))
|
||||
|
||||
# Test shape of serialized embeddings
|
||||
with open(stream, "rb") as queue:
|
||||
self.assertEqual(SparseArray().load(queue).shape, (1, 30522))
|
||||
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
Sparse Vectors module tests
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from txtai.vectors import SparseVectors, SparseVectorsFactory
|
||||
|
||||
|
||||
class TestSparseVectors(unittest.TestCase):
|
||||
"""
|
||||
Sparse Vectors tests.
|
||||
"""
|
||||
|
||||
def testCustom(self):
|
||||
"""
|
||||
Test custom sparse vectors instance
|
||||
"""
|
||||
|
||||
self.assertIsNotNone(
|
||||
SparseVectorsFactory.create({"method": "txtai.vectors.SparseSTVectors", "path": "sparse-encoder-testing/splade-bert-tiny-nq"})
|
||||
)
|
||||
|
||||
def testDefaultNormalize(self):
|
||||
"""
|
||||
Test defaultnormalize method
|
||||
"""
|
||||
|
||||
vectors = SparseVectors(None, None, None)
|
||||
self.assertFalse(vectors.defaultnormalize())
|
||||
|
||||
def testNotSupported(self):
|
||||
"""
|
||||
Test exceptions for unsupported methods
|
||||
"""
|
||||
|
||||
vectors = SparseVectors(None, None, None)
|
||||
|
||||
self.assertRaises(ValueError, vectors.truncate, None)
|
||||
self.assertRaises(ValueError, vectors.quantize, None)
|
||||
|
||||
def testNotFound(self):
|
||||
"""
|
||||
Test unresolvable vector backend
|
||||
"""
|
||||
|
||||
with self.assertRaises(ImportError):
|
||||
SparseVectorsFactory.create({"method": "notfound.vectors"})
|
||||
@@ -0,0 +1,609 @@
|
||||
"""
|
||||
Workflow module tests
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import glob
|
||||
import io
|
||||
import os
|
||||
import tempfile
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from txtai.api import API
|
||||
from txtai.embeddings import Documents, Embeddings
|
||||
from txtai.pipeline import Nop, Segmentation, Summary, Translation, Textractor
|
||||
from txtai.workflow import (
|
||||
Workflow,
|
||||
Task,
|
||||
ConsoleTask,
|
||||
ExportTask,
|
||||
FileTask,
|
||||
ImageTask,
|
||||
RagTask,
|
||||
RetrieveTask,
|
||||
StorageTask,
|
||||
TemplateTask,
|
||||
WorkflowTask,
|
||||
)
|
||||
|
||||
# pylint: disable=C0411
|
||||
from utils import Utils
|
||||
|
||||
|
||||
# pylint: disable=R0904
|
||||
class TestWorkflow(unittest.TestCase):
|
||||
"""
|
||||
Workflow tests.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""
|
||||
Initialize test data.
|
||||
"""
|
||||
|
||||
# Default YAML workflow configuration
|
||||
cls.config = """
|
||||
# Embeddings index
|
||||
writable: true
|
||||
embeddings:
|
||||
scoring: bm25
|
||||
path: google/bert_uncased_L-2_H-128_A-2
|
||||
content: true
|
||||
|
||||
# Text segmentation
|
||||
segmentation:
|
||||
sentences: true
|
||||
|
||||
# Workflow definitions
|
||||
workflow:
|
||||
index:
|
||||
tasks:
|
||||
- action: segmentation
|
||||
- action: index
|
||||
search:
|
||||
tasks:
|
||||
- search
|
||||
transform:
|
||||
tasks:
|
||||
- transform
|
||||
"""
|
||||
|
||||
def testBaseWorkflow(self):
|
||||
"""
|
||||
Test a basic workflow
|
||||
"""
|
||||
|
||||
translate = Translation()
|
||||
|
||||
# Workflow that translate text to Spanish
|
||||
workflow = Workflow([Task(lambda x: translate(x, "es"))])
|
||||
|
||||
results = list(workflow(["The sky is blue", "Forest through the trees"]))
|
||||
|
||||
self.assertEqual(len(results), 2)
|
||||
|
||||
def testChainWorkflow(self):
|
||||
"""
|
||||
Test a chain of workflows
|
||||
"""
|
||||
|
||||
workflow1 = Workflow([Task(lambda x: [y * 2 for y in x])])
|
||||
workflow2 = Workflow([Task(lambda x: [y - 1 for y in x])], batch=4)
|
||||
|
||||
results = list(workflow2(workflow1([1, 2, 4, 8, 16, 32])))
|
||||
self.assertEqual(results, [1, 3, 7, 15, 31, 63])
|
||||
|
||||
def testComplexWorkflow(self):
|
||||
"""
|
||||
Test a complex workflow
|
||||
"""
|
||||
|
||||
textractor = Textractor(paragraphs=True, minlength=150, join=True)
|
||||
summary = Summary("t5-small")
|
||||
|
||||
embeddings = Embeddings({"path": "sentence-transformers/nli-mpnet-base-v2"})
|
||||
documents = Documents()
|
||||
|
||||
def index(x):
|
||||
documents.add(x)
|
||||
return x
|
||||
|
||||
# Extract text and summarize articles
|
||||
articles = Workflow([FileTask(textractor), Task(lambda x: summary(x, maxlength=15))])
|
||||
|
||||
# Complex workflow that extracts text, runs summarization then loads into an embeddings index
|
||||
tasks = [WorkflowTask(articles, r".\.pdf$"), Task(index, unpack=False)]
|
||||
|
||||
data = ["file://" + Utils.PATH + "/article.pdf", "Workflows can process audio files, documents and snippets"]
|
||||
|
||||
# Convert file paths to data tuples
|
||||
data = [(x, element, None) for x, element in enumerate(data)]
|
||||
|
||||
# Execute workflow, discard results as they are streamed
|
||||
workflow = Workflow(tasks)
|
||||
data = list(workflow(data))
|
||||
|
||||
# Build the embeddings index
|
||||
embeddings.index(documents)
|
||||
|
||||
# Cleanup temporary storage
|
||||
documents.close()
|
||||
|
||||
# Run search and validate result
|
||||
index, _ = embeddings.search("search text", 1)[0]
|
||||
self.assertEqual(index, 0)
|
||||
self.assertEqual(data[0][1], "txtai builds an AI-powered index over sections")
|
||||
|
||||
def testConcurrentWorkflow(self):
|
||||
"""
|
||||
Test running concurrent task actions
|
||||
"""
|
||||
|
||||
nop = Nop()
|
||||
|
||||
workflow = Workflow([Task([nop, nop], concurrency="thread")])
|
||||
results = list(workflow([2, 4]))
|
||||
self.assertEqual(results, [(2, 2), (4, 4)])
|
||||
|
||||
workflow = Workflow([Task([nop, nop], concurrency="process")])
|
||||
results = list(workflow([2, 4]))
|
||||
self.assertEqual(results, [(2, 2), (4, 4)])
|
||||
|
||||
workflow = Workflow([Task([nop, nop], concurrency="unknown")])
|
||||
results = list(workflow([2, 4]))
|
||||
self.assertEqual(results, [(2, 2), (4, 4)])
|
||||
|
||||
def testConsoleWorkflow(self):
|
||||
"""
|
||||
Test a console task
|
||||
"""
|
||||
|
||||
# Excel export
|
||||
workflow = Workflow([ConsoleTask()])
|
||||
|
||||
output = io.StringIO()
|
||||
with contextlib.redirect_stdout(output):
|
||||
list(workflow([{"id": 1, "text": "Sentence 1"}, {"id": 2, "text": "Sentence 2"}]))
|
||||
|
||||
self.assertIn("Sentence 2", output.getvalue())
|
||||
|
||||
def testExportWorkflow(self):
|
||||
"""
|
||||
Test an export task
|
||||
"""
|
||||
|
||||
# Excel export
|
||||
path = os.path.join(tempfile.gettempdir(), "export.xlsx")
|
||||
workflow = Workflow([ExportTask(output=path)])
|
||||
list(workflow([{"id": 1, "text": "Sentence 1"}, {"id": 2, "text": "Sentence 2"}]))
|
||||
self.assertGreater(os.path.getsize(path), 0)
|
||||
|
||||
# Export CSV
|
||||
path = os.path.join(tempfile.gettempdir(), "export.csv")
|
||||
workflow = Workflow([ExportTask(output=path)])
|
||||
list(workflow([{"id": 1, "text": "Sentence 1"}, {"id": 2, "text": "Sentence 2"}]))
|
||||
self.assertGreater(os.path.getsize(path), 0)
|
||||
|
||||
# Export CSV with timestamp
|
||||
path = os.path.join(tempfile.gettempdir(), "export-timestamp.csv")
|
||||
workflow = Workflow([ExportTask(output=path, timestamp=True)])
|
||||
list(workflow([{"id": 1, "text": "Sentence 1"}, {"id": 2, "text": "Sentence 2"}]))
|
||||
|
||||
# Find timestamped file and ensure it has data
|
||||
path = glob.glob(os.path.join(tempfile.gettempdir(), "export-timestamp*.csv"))[0]
|
||||
self.assertGreater(os.path.getsize(path), 0)
|
||||
|
||||
def testExtractWorkflow(self):
|
||||
"""
|
||||
Test column extraction tasks
|
||||
"""
|
||||
|
||||
workflow = Workflow([Task(lambda x: x, unpack=False, column=0)], batch=1)
|
||||
|
||||
results = list(workflow([(0, 1)]))
|
||||
self.assertEqual(results[0], 0)
|
||||
|
||||
results = list(workflow([(0, (1, 2), None)]))
|
||||
self.assertEqual(results[0], (0, 1, None))
|
||||
|
||||
results = list(workflow([1]))
|
||||
self.assertEqual(results[0], 1)
|
||||
|
||||
def testImageWorkflow(self):
|
||||
"""
|
||||
Test an image task
|
||||
"""
|
||||
|
||||
workflow = Workflow([ImageTask()])
|
||||
|
||||
results = list(workflow([Utils.PATH + "/books.jpg"]))
|
||||
|
||||
self.assertEqual(results[0].size, (1024, 682))
|
||||
|
||||
def testInvalidWorkflow(self):
|
||||
"""
|
||||
Test task with invalid parameters
|
||||
"""
|
||||
|
||||
with self.assertRaises(TypeError):
|
||||
Task(invalid=True)
|
||||
|
||||
def testMergeWorkflow(self):
|
||||
"""
|
||||
Test merge tasks
|
||||
"""
|
||||
|
||||
task = Task([lambda x: [pow(y, 2) for y in x], lambda x: [pow(y, 3) for y in x]], merge="hstack")
|
||||
|
||||
# Test hstack (column-wise) merge
|
||||
workflow = Workflow([task])
|
||||
results = list(workflow([2, 4]))
|
||||
self.assertEqual(results, [(4, 8), (16, 64)])
|
||||
|
||||
# Test vstack (row-wise) merge
|
||||
task.merge = "vstack"
|
||||
results = list(workflow([2, 4]))
|
||||
self.assertEqual(results, [4, 8, 16, 64])
|
||||
|
||||
# Test concat (values joined into single string) merge
|
||||
task.merge = "concat"
|
||||
results = list(workflow([2, 4]))
|
||||
self.assertEqual(results, ["4. 8", "16. 64"])
|
||||
|
||||
# Test no merge
|
||||
task.merge = None
|
||||
results = list(workflow([2, 4, 6]))
|
||||
self.assertEqual(results, [[4, 16, 36], [8, 64, 216]])
|
||||
|
||||
# Test generated (id, data, tag) tuples are properly returned
|
||||
workflow = Workflow([Task(lambda x: [(0, y, None) for y in x])])
|
||||
results = list(workflow([(1, "text", "tags")]))
|
||||
self.assertEqual(results[0], (0, "text", None))
|
||||
|
||||
def testMergeUnbalancedWorkflow(self):
|
||||
"""
|
||||
Test merge tasks with unbalanced outputs (i.e. one action produce more output than another for same input).
|
||||
"""
|
||||
|
||||
nop = Nop()
|
||||
segment1 = Segmentation(sentences=True)
|
||||
|
||||
task = Task([nop, segment1])
|
||||
|
||||
# Test hstack
|
||||
workflow = Workflow([task])
|
||||
results = list(workflow(["This is a test sentence. And another sentence to split."]))
|
||||
self.assertEqual(
|
||||
results, [("This is a test sentence. And another sentence to split.", ["This is a test sentence.", "And another sentence to split."])]
|
||||
)
|
||||
|
||||
# Test vstack
|
||||
task.merge = "vstack"
|
||||
workflow = Workflow([task])
|
||||
results = list(workflow(["This is a test sentence. And another sentence to split."]))
|
||||
self.assertEqual(
|
||||
results, ["This is a test sentence. And another sentence to split.", "This is a test sentence.", "And another sentence to split."]
|
||||
)
|
||||
|
||||
def testNumpyWorkflow(self):
|
||||
"""
|
||||
Test a numpy workflow
|
||||
"""
|
||||
|
||||
task = Task([lambda x: np.power(x, 2), lambda x: np.power(x, 3)], merge="hstack")
|
||||
|
||||
# Test hstack (column-wise) merge
|
||||
workflow = Workflow([task])
|
||||
results = list(workflow(np.array([2, 4])))
|
||||
self.assertTrue(np.array_equal(np.array(results), np.array([[4, 8], [16, 64]])))
|
||||
|
||||
# Test vstack (row-wise) merge
|
||||
task.merge = "vstack"
|
||||
results = list(workflow(np.array([2, 4])))
|
||||
self.assertEqual(results, [4, 8, 16, 64])
|
||||
|
||||
# Test no merge
|
||||
task.merge = None
|
||||
results = list(workflow(np.array([2, 4, 6])))
|
||||
self.assertTrue(np.array_equal(np.array(results), np.array([[4, 16, 36], [8, 64, 216]])))
|
||||
|
||||
def testRetrieveWorkflow(self):
|
||||
"""
|
||||
Test a retrieve task
|
||||
"""
|
||||
|
||||
# Test retrieve with generated temporary directory
|
||||
workflow = Workflow([RetrieveTask()])
|
||||
results = list(workflow(["file://" + Utils.PATH + "/books.jpg"]))
|
||||
self.assertTrue(results[0].endswith("books.jpg"))
|
||||
|
||||
# Test retrieve with specified temporary directory
|
||||
workflow = Workflow([RetrieveTask(directory=os.path.join(tempfile.gettempdir(), "retrieve"))])
|
||||
results = list(workflow(["file://" + Utils.PATH + "/books.jpg"]))
|
||||
self.assertTrue(results[0].endswith("books.jpg"))
|
||||
|
||||
# Test with directory structures
|
||||
workflow = Workflow([RetrieveTask(flatten=False)])
|
||||
results = list(workflow(["file://" + Utils.PATH + "/books.jpg"]))
|
||||
self.assertTrue(results[0].endswith("books.jpg") and "txtai" in results[0])
|
||||
|
||||
def testScheduleWorkflow(self):
|
||||
"""
|
||||
Test workflow schedules
|
||||
"""
|
||||
|
||||
# Test workflow schedule with Python
|
||||
workflow = Workflow([Task()])
|
||||
workflow.schedule("* * * * * *", ["test"], 1)
|
||||
self.assertEqual(len(workflow.tasks), 1)
|
||||
|
||||
# Test workflow schedule with YAML
|
||||
workflow = """
|
||||
segmentation:
|
||||
sentences: true
|
||||
workflow:
|
||||
segment:
|
||||
schedule:
|
||||
cron: '* * * * * *'
|
||||
elements:
|
||||
- a sentence to segment
|
||||
iterations: 1
|
||||
tasks:
|
||||
- action: segmentation
|
||||
task: console
|
||||
"""
|
||||
|
||||
output = io.StringIO()
|
||||
with contextlib.redirect_stdout(output):
|
||||
app = API(workflow)
|
||||
app.wait()
|
||||
|
||||
self.assertIn("a sentence to segment", output.getvalue())
|
||||
|
||||
def testScheduleErrorWorkflow(self):
|
||||
"""
|
||||
Test workflow schedules with errors
|
||||
"""
|
||||
|
||||
def action(elements):
|
||||
raise FileNotFoundError
|
||||
|
||||
# Test workflow proceeds after exception raised
|
||||
with self.assertLogs() as logs:
|
||||
workflow = Workflow([Task(action=action)])
|
||||
workflow.schedule("* * * * * *", ["test"], 1)
|
||||
|
||||
self.assertIn("FileNotFoundError", " ".join(logs.output))
|
||||
|
||||
def testStorageWorkflow(self):
|
||||
"""
|
||||
Test a storage task
|
||||
"""
|
||||
|
||||
workflow = Workflow([StorageTask()])
|
||||
|
||||
results = list(workflow(["local://" + Utils.PATH, "test string"]))
|
||||
|
||||
self.assertEqual(len(results), 22)
|
||||
|
||||
def testTemplateInput(self):
|
||||
"""
|
||||
Test template task input
|
||||
"""
|
||||
|
||||
workflow = Workflow([TemplateTask(template="This is a {text}")])
|
||||
|
||||
# Test with string inputs
|
||||
results = list(workflow(["prompt"]))
|
||||
self.assertEqual(results[0], "This is a prompt")
|
||||
|
||||
# Test with dict inputs
|
||||
results = list(workflow([{"text": "prompt"}]))
|
||||
self.assertEqual(results[0], "This is a prompt")
|
||||
|
||||
# Test with tuple inputs
|
||||
workflow = Workflow([TemplateTask(template="This is a {arg0}", unpack=False)])
|
||||
results = list(workflow([("prompt",)]))
|
||||
self.assertEqual(results[0], "This is a prompt")
|
||||
|
||||
# Test invalid inputs
|
||||
with self.assertRaises(KeyError):
|
||||
workflow = Workflow([TemplateTask(template="No variables")])
|
||||
results = list(workflow([{"unused": "prompt"}]))
|
||||
|
||||
# Test no template
|
||||
workflow = Workflow([TemplateTask()])
|
||||
results = list(workflow(["prompt"]))
|
||||
self.assertEqual(results[0], "prompt")
|
||||
|
||||
def testTemplateRules(self):
|
||||
"""
|
||||
Test template task rules
|
||||
"""
|
||||
|
||||
# Test rule applied
|
||||
workflow = Workflow([TemplateTask(template="This is a {text}", rules={"text": "Test skip"})])
|
||||
results = list(workflow([{"text": "Test skip"}]))
|
||||
self.assertEqual(results[0], "Test skip")
|
||||
|
||||
# Test rule not applied
|
||||
results = list(workflow([{"text": "prompt"}]))
|
||||
self.assertEqual(results[0], "This is a prompt")
|
||||
|
||||
def testTemplateRag(self):
|
||||
"""
|
||||
Test rag template task
|
||||
"""
|
||||
|
||||
# Test outputs
|
||||
workflow = Workflow([RagTask(template="This is a {text}")])
|
||||
results = list(workflow(["prompt"]))
|
||||
self.assertEqual(results[0], {"query": "prompt", "question": "This is a prompt"})
|
||||
|
||||
# Test partial outputs
|
||||
workflow = Workflow([RagTask(template="This is a {text}")])
|
||||
results = list(workflow([{"query": "query", "question": "prompt"}]))
|
||||
self.assertEqual(results[0], {"query": "query", "question": "This is a prompt"})
|
||||
|
||||
# Test additional template parameters
|
||||
workflow = Workflow([RagTask(template="This is a {text} with another {param}")])
|
||||
results = list(workflow([{"query": "query", "question": "prompt", "param": "value"}]))
|
||||
self.assertEqual(results[0], {"query": "query", "question": "This is a prompt with another value", "param": "value"})
|
||||
|
||||
def testTensorTransformWorkflow(self):
|
||||
"""
|
||||
Test a tensor workflow with list transformations
|
||||
"""
|
||||
|
||||
# Test one-one list transformation
|
||||
task = Task(lambda x: x.tolist())
|
||||
workflow = Workflow([task])
|
||||
results = list(workflow(np.array([2])))
|
||||
self.assertEqual(results, [2])
|
||||
|
||||
# Test one-many list transformation
|
||||
task = Task(lambda x: [x.tolist() * 2])
|
||||
workflow = Workflow([task])
|
||||
results = list(workflow(np.array([2])))
|
||||
self.assertEqual(results, [2, 2])
|
||||
|
||||
def testTorchWorkflow(self):
|
||||
"""
|
||||
Test a torch workflow
|
||||
"""
|
||||
|
||||
# pylint: disable=E1101,E1102
|
||||
task = Task([lambda x: torch.pow(x, 2), lambda x: torch.pow(x, 3)], merge="hstack")
|
||||
|
||||
# Test hstack (column-wise) merge
|
||||
workflow = Workflow([task])
|
||||
results = np.array([x.numpy() for x in workflow(torch.tensor([2, 4]))])
|
||||
self.assertTrue(np.array_equal(results, np.array([[4, 8], [16, 64]])))
|
||||
|
||||
# Test vstack (row-wise) merge
|
||||
task.merge = "vstack"
|
||||
results = list(workflow(torch.tensor([2, 4])))
|
||||
self.assertEqual(results, [4, 8, 16, 64])
|
||||
|
||||
# Test no merge
|
||||
task.merge = None
|
||||
results = np.array([x.numpy() for x in workflow(torch.tensor([2, 4, 6]))])
|
||||
self.assertTrue(np.array_equal(np.array(results), np.array([[4, 16, 36], [8, 64, 216]])))
|
||||
|
||||
def testYamlFunctionWorkflow(self):
|
||||
"""
|
||||
Test YAML workflow with a function action
|
||||
"""
|
||||
|
||||
# Create function and add to module
|
||||
def action(elements):
|
||||
return [x * 2 for x in elements]
|
||||
|
||||
sys.modules[__name__].action = action
|
||||
|
||||
workflow = """
|
||||
workflow:
|
||||
run:
|
||||
tasks:
|
||||
- testworkflow.action
|
||||
"""
|
||||
|
||||
app = API(workflow)
|
||||
self.assertEqual(list(app.workflow("run", [1, 2])), [2, 4])
|
||||
|
||||
def testYamlIndexWorkflow(self):
|
||||
"""
|
||||
Test reading a YAML index workflow in Python.
|
||||
"""
|
||||
|
||||
app = API(self.config)
|
||||
self.assertEqual(
|
||||
list(app.workflow("index", ["This is a test sentence. And another sentence to split."])),
|
||||
["This is a test sentence.", "And another sentence to split."],
|
||||
)
|
||||
|
||||
# Read from file
|
||||
path = os.path.join(tempfile.gettempdir(), "workflow.yml")
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write(self.config)
|
||||
|
||||
app = API(path)
|
||||
self.assertEqual(
|
||||
list(app.workflow("index", ["This is a test sentence. And another sentence to split."])),
|
||||
["This is a test sentence.", "And another sentence to split."],
|
||||
)
|
||||
|
||||
# Read from YAML object
|
||||
app = API(API.read(self.config))
|
||||
self.assertEqual(
|
||||
list(app.workflow("index", ["This is a test sentence. And another sentence to split."])),
|
||||
["This is a test sentence.", "And another sentence to split."],
|
||||
)
|
||||
|
||||
def testYamlSearchWorkflow(self):
|
||||
"""
|
||||
Test reading a YAML search workflow in Python.
|
||||
"""
|
||||
|
||||
# Test search
|
||||
app = API(self.config)
|
||||
list(app.workflow("index", ["This is a test sentence. And another sentence to split."]))
|
||||
self.assertEqual(
|
||||
list(app.workflow("search", ["another"]))[0]["text"],
|
||||
"And another sentence to split.",
|
||||
)
|
||||
|
||||
def testYamlWorkflowTask(self):
|
||||
"""
|
||||
Test YAML workflow with a workflow task
|
||||
"""
|
||||
|
||||
# Create function and add to module
|
||||
def action(elements):
|
||||
return [x * 2 for x in elements]
|
||||
|
||||
sys.modules[__name__].action = action
|
||||
|
||||
workflow = """
|
||||
workflow:
|
||||
run:
|
||||
tasks:
|
||||
- testworkflow.action
|
||||
flow:
|
||||
tasks:
|
||||
- run
|
||||
"""
|
||||
|
||||
app = API(workflow)
|
||||
self.assertEqual(list(app.workflow("flow", [1, 2])), [2, 4])
|
||||
|
||||
def testYamlTransformWorkflow(self):
|
||||
"""
|
||||
Test reading a YAML transform workflow in Python.
|
||||
"""
|
||||
|
||||
# Test search
|
||||
app = API(self.config)
|
||||
self.assertEqual(len(list(app.workflow("transform", ["text"]))[0]), 128)
|
||||
|
||||
def testYamlError(self):
|
||||
"""
|
||||
Test reading a YAML workflow with errors.
|
||||
"""
|
||||
|
||||
# Read from string
|
||||
config = """
|
||||
# Workflow definitions
|
||||
workflow:
|
||||
error:
|
||||
tasks:
|
||||
- action: error
|
||||
"""
|
||||
|
||||
with self.assertRaises(KeyError):
|
||||
API(config)
|
||||
@@ -0,0 +1,11 @@
|
||||
"""
|
||||
Utils module
|
||||
"""
|
||||
|
||||
|
||||
class Utils:
|
||||
"""
|
||||
Utility constants and methods
|
||||
"""
|
||||
|
||||
PATH = "/tmp/txtai"
|
||||
Reference in New Issue
Block a user