chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
Reference in New Issue
Block a user