feat: Implementation of Table Cell Proposal (#4616)

* Starting adding support for TableCell

* Update tests to use row and col

* Added schema test to check to_dict and from_dict works for Table documents. Also updated Doc.__eq__ to work for tables.

* Update eval test to use TableCell

* Added more schema tests for table docs, labels and answers.

* Add boolean to toggle between Span and TableCell

* Add deprecation message

* Test that table answers work as responses in the rest API

---------

Co-authored-by: agnieszka-m <[email protected]>
This commit is contained in:
Sebastian
2023-04-19 13:14:49 +02:00
committed by GitHub
co-authored by agnieszka-m
parent ec7fc4aa0b
commit 8d9136bad4
10 changed files with 373 additions and 143 deletions
+1 -1
View File
@@ -17,7 +17,7 @@ import logging
import pandas as pd
from haystack.schema import Document, Answer, Label, MultiLabel, Span, EvaluationResult
from haystack.schema import Document, Answer, Label, MultiLabel, Span, EvaluationResult, TableCell
from haystack.nodes.base import BaseComponent
from haystack.pipelines.base import Pipeline
from haystack.environment import set_pytorch_secure_model_loading
+8 -2
View File
@@ -1058,6 +1058,12 @@ class FARMReader(BaseReader):
if not label.document.id:
logger.error("Label does not contain a document id")
continue
if label.document.content_type == "table":
logger.warning(
"Label with a table document is not compatible with the FARMReader. Skipping label with id %s.",
label.id,
)
continue
aggregated_per_doc[label.document.id].append(label)
# Create squad style dicts
@@ -1101,7 +1107,7 @@ class FARMReader(BaseReader):
)
continue
aggregated_per_question[aggregation_key]["answers"].append(
{"text": label.answer.answer, "answer_start": label.answer.offsets_in_document[0].start}
{"text": label.answer.answer, "answer_start": label.answer.offsets_in_document[0].start} # type: ignore [union-attr]
)
aggregated_per_question[aggregation_key]["is_impossible"] = False
# create new one
@@ -1121,7 +1127,7 @@ class FARMReader(BaseReader):
"answers": [
{
"text": label.answer.answer,
"answer_start": label.answer.offsets_in_document[0].start,
"answer_start": label.answer.offsets_in_document[0].start, # type: ignore [union-attr]
}
],
"is_impossible": False,
+80 -10
View File
@@ -25,7 +25,7 @@ from transformers import (
from transformers.models.tapas.modeling_tapas import TapasPreTrainedModel
from haystack.errors import HaystackError
from haystack.schema import Document, Answer, Span
from haystack.schema import Document, Answer, TableCell, Span
from haystack.nodes.reader.base import BaseReader
from haystack.modeling.utils import initialize_device_settings
@@ -71,6 +71,7 @@ class TableReader(BaseReader):
max_seq_len: int = 256,
use_auth_token: Optional[Union[str, bool]] = None,
devices: Optional[List[Union[str, torch.device]]] = None,
return_table_cell: bool = False,
):
"""
Load a TableQA model from Transformers.
@@ -112,9 +113,23 @@ class TableReader(BaseReader):
A list containing torch device objects and/or strings is supported (For example
[torch.device('cuda:0'), "mps", "cuda:1"]). When specifying `use_gpu=False` the devices
parameter is not used and a single cpu device is used for inference.
:param return_table_cell: Whether to return the offsets (`offsets_in_document` and `offsets_in_context`) that indicate
the table cells that answer the question using the TableCell schema. The TableCell schema returns the row
and column indices of the table cells selected in the Answer. Otherwise, the offsets
are returned as Span objects which are start and end indices when counting through the
table in a linear fashion, which means the first cell is top left and the last cell is bottom right.
"""
super().__init__()
if not return_table_cell:
logger.warning(
"The support for returning offsets in answer predictions in a linear fashion is being deprecated."
" Set return_table_cell=True to use the new offsets format which returns the row and column indices"
" of the table cells selected in the answer."
" In the future, return_table_cell=True will become default and return_table_cell=False will no "
" longer be supported."
)
self.devices, _ = initialize_device_settings(devices=devices, use_cuda=use_gpu, multi_gpu=False)
if len(self.devices) > 1:
logger.warning(
@@ -133,6 +148,7 @@ class TableReader(BaseReader):
tokenizer=tokenizer,
max_seq_len=max_seq_len,
use_auth_token=use_auth_token,
return_table_cell=return_table_cell,
)
elif config.architectures[0] == "TapasForScoredQA":
self.table_encoder = _TapasScoredEncoder(
@@ -144,6 +160,7 @@ class TableReader(BaseReader):
return_no_answer=return_no_answer,
max_seq_len=max_seq_len,
use_auth_token=use_auth_token,
return_table_cell=return_table_cell,
)
else:
logger.error(
@@ -238,6 +255,7 @@ class _TapasEncoder:
tokenizer: Optional[str] = None,
max_seq_len: int = 256,
use_auth_token: Optional[Union[str, bool]] = None,
return_table_cell: bool = False,
):
self.model = TapasForQuestionAnswering.from_pretrained(
model_name_or_path, revision=model_version, use_auth_token=use_auth_token
@@ -259,6 +277,7 @@ class _TapasEncoder:
framework="pt",
batch_size=1, # batch_size of 1 only works currently b/c of issue with HuggingFace pipeline logic and the return type of TableQuestionAnsweringPipeline._forward
device=self.device,
return_table_cell=return_table_cell,
)
def predict(self, query: str, documents: List[Document], top_k: int) -> Dict:
@@ -300,6 +319,10 @@ class _TapasEncoder:
class _TableQuestionAnsweringPipeline(TableQuestionAnsweringPipeline):
"""Modified from transformers TableQuestionAnsweringPipeline.postprocess to return Haystack Answer objects."""
def __init__(self, *args, return_table_cell: bool = False, **kwargs):
super().__init__(*args, **kwargs)
self.return_table_cell = return_table_cell
def _calculate_answer_score(
self, logits: torch.Tensor, inputs: Dict, answer_coordinates: List[List[Tuple[int, int]]]
) -> List[Optional[np.ndarray]]:
@@ -413,8 +436,11 @@ class _TableQuestionAnsweringPipeline(TableQuestionAnsweringPipeline):
answer_str = ", ".join(cells)
else:
answer_str = self._aggregate_answers(aggregator, cells)
answer_offsets = _calculate_answer_offsets(ans_coordinates_per_table, string_table)
current_score = answer_scores[index]
if self.return_table_cell:
answer_offsets = _calculate_answer_offsets(ans_coordinates_per_table)
else:
answer_offsets = _calculate_answer_offsets_span(ans_coordinates_per_table, string_table)
answer = Answer(
answer=answer_str,
type="extractive",
@@ -445,6 +471,7 @@ class _TapasScoredEncoder:
return_no_answer: bool = False,
max_seq_len: int = 256,
use_auth_token: Optional[Union[str, bool]] = None,
return_table_cell: bool = False,
):
self.model = self._TapasForScoredQA.from_pretrained(
model_name_or_path, revision=model_version, use_auth_token=use_auth_token
@@ -457,6 +484,7 @@ class _TapasScoredEncoder:
self.device = device
self.top_k_per_candidate = top_k_per_candidate
self.return_no_answer = return_no_answer
self.return_table_cell = return_table_cell
def _predict_tapas_scored(self, inputs: BatchEncoding, document: Document) -> Tuple[List[Answer], float]:
orig_table: pd.DataFrame = document.content
@@ -537,7 +565,11 @@ class _TapasScoredEncoder:
for answer_span_idx in top_k_answer_spans.indices:
current_answer_span = possible_answer_spans[answer_span_idx]
answer_str = string_table.iat[current_answer_span[:2]]
answer_offsets = _calculate_answer_offsets([current_answer_span[:2]], string_table)
answer_offsets: Union[List[Span], List[TableCell]]
if self.return_table_cell:
answer_offsets = _calculate_answer_offsets([current_answer_span[:2]])
else:
answer_offsets = _calculate_answer_offsets_span([current_answer_span[:2]], string_table)
# As the general table score is more important for the final score, it is double weighted.
current_score = ((2 * table_relevancy_prob) + span_logits_softmax[0, answer_span_idx].item()) / 3
@@ -574,14 +606,20 @@ class _TapasScoredEncoder:
no_answer_score = current_no_answer_score
if self.return_no_answer:
if self.return_table_cell:
offsets_in_context = None
offsets_in_document = None
else:
offsets_in_context = [Span(start=0, end=0)]
offsets_in_document = [Span(start=0, end=0)]
answers.append(
Answer(
answer="",
type="extractive",
score=no_answer_score,
context=None,
offsets_in_context=[Span(start=0, end=0)],
offsets_in_document=[Span(start=0, end=0)],
offsets_in_context=offsets_in_context,
offsets_in_document=offsets_in_document,
document_ids=None,
meta=None,
)
@@ -649,6 +687,7 @@ class RCIReader(BaseReader):
top_k: int = 10,
max_seq_len: int = 256,
use_auth_token: Optional[Union[str, bool]] = None,
return_table_cell: bool = False,
):
"""
Load an RCI model from Transformers.
@@ -677,9 +716,23 @@ class RCIReader(BaseReader):
`transformers-cli login` (stored in ~/.huggingface) will be used.
Additional information can be found here
https://huggingface.co/transformers/main_classes/model.html#transformers.PreTrainedModel.from_pretrained
:param return_table_cell: Whether to return the offsets (`offsets_in_document` and `offsets_in_context`) that indicate
the cells that answer the question using the TableCell schema. The TableCell schema returns the row
and column indices of the table cells selected in the Answer. Otherwise, the offsets
are returned as Span objects which are start and end indices when counting through the
table in a linear fashion, which means the first cell is top left and the last cell is bottom right.
"""
super().__init__()
if not return_table_cell:
logger.warning(
"The support for returning offsets in answer predictions in a linear fashion is being deprecated."
" Set return_table_cell=True to use the new offsets format which returns the row and column indices"
" of the table cells selected in the answer."
" In the future, return_table_cell=True will become default and return_table_cell=False will no "
" longer be supported."
)
self.devices, _ = initialize_device_settings(use_cuda=use_gpu, multi_gpu=False)
if len(self.devices) > 1:
logger.warning(
@@ -722,6 +775,7 @@ class RCIReader(BaseReader):
self.top_k = top_k
self.max_seq_len = max_seq_len
self.return_no_answers = False
self.return_table_cell = return_table_cell
def predict(self, query: str, documents: List[Document], top_k: Optional[int] = None) -> Dict:
"""
@@ -787,15 +841,19 @@ class RCIReader(BaseReader):
cell_scores_table[-1].append(current_cell_score)
answer_str = string_table.iloc[row_idx, col_idx]
answer_offsets = self._calculate_answer_offsets(row_idx, col_idx, string_table)
answer_offsets: Union[List[Span], List[TableCell]]
if self.return_table_cell:
answer_offsets = [TableCell(row=row_idx, col=col_idx)]
else:
answer_offsets = [self._calculate_answer_offsets_span(row_idx, col_idx, string_table)]
current_answers.append(
Answer(
answer=answer_str,
type="extractive",
score=current_cell_score,
context=string_table,
offsets_in_document=[answer_offsets],
offsets_in_context=[answer_offsets],
offsets_in_document=answer_offsets,
offsets_in_context=answer_offsets,
document_ids=[document.id],
)
)
@@ -831,7 +889,7 @@ class RCIReader(BaseReader):
return row_reps, column_reps
@staticmethod
def _calculate_answer_offsets(row_idx, column_index, table) -> Span:
def _calculate_answer_offsets_span(row_idx, column_index, table) -> Span:
_, n_columns = table.shape
answer_cell_offset = (row_idx * n_columns) + column_index
@@ -868,7 +926,19 @@ class RCIReader(BaseReader):
return results
def _calculate_answer_offsets(answer_coordinates: List[Tuple[int, int]], table: pd.DataFrame) -> List[Span]:
def _calculate_answer_offsets(answer_coordinates: List[Tuple[int, int]]) -> List[TableCell]:
"""
Calculates the answer cell offsets of the linearized table based on the answer cell coordinates.
:param answer_coordinates: List of answer coordinates.
"""
answer_offsets = []
for coord in answer_coordinates:
answer_offsets.append(TableCell(row=coord[0], col=coord[1]))
return answer_offsets
def _calculate_answer_offsets_span(answer_coordinates: List[Tuple[int, int]], table: pd.DataFrame) -> List[Span]:
"""
Calculates the answer cell offsets of the linearized table based on the answer cell coordinates.
+10 -4
View File
@@ -1621,8 +1621,10 @@ class Pipeline:
df_answers["gold_contexts_similarity"] = df_answers.map_rows(
lambda row: [
calculate_context_similarity(
str(gold_context), # could be dataframe
str(row["context"]) if row["context"] is not None else "", # could be dataframe
str(gold_context), # could be dataframe or list of lists
str(row["context"])
if row["context"] is not None
else "", # could be dataframe or list of lists
min_length=context_matching_min_length,
boost_split_overlaps=context_matching_boost_split_overlaps,
)
@@ -1759,9 +1761,13 @@ class Pipeline:
df_docs["gold_contexts_similarity"] = df_docs.map_rows(
lambda row: [
calculate_context_similarity(
str(gold_context) if isinstance(gold_context, pd.DataFrame) else gold_context,
str(gold_context)
if isinstance(gold_context, (pd.DataFrame, list))
else gold_context, # could be dataframe or list of lists
str(row["context"])
if isinstance(row["context"], pd.DataFrame)
if isinstance(
row["context"], (pd.DataFrame, list)
) # could be dataframe or list of lists
else row["context"] or "",
min_length=context_matching_min_length,
boost_split_overlaps=context_matching_boost_split_overlaps,
+47 -10
View File
@@ -250,9 +250,14 @@ class Document:
return cls.from_dict(dictionary, field_map=field_map)
def __eq__(self, other):
content = getattr(other, "content", None)
if isinstance(content, pd.DataFrame):
is_content_equal = content.equals(self.content)
else:
is_content_equal = content == self.content
return (
isinstance(other, self.__class__)
and getattr(other, "content", None) == self.content
and is_content_equal
and getattr(other, "content_type", None) == self.content_type
and getattr(other, "id", None) == self.id
and getattr(other, "id_hash_keys", None) == self.id_hash_keys
@@ -286,10 +291,9 @@ class Span:
"""
Defining a sequence of characters (Text span) or cells (Table span) via start and end index.
For extractive QA: Character where answer starts/ends
For TableQA: Cell where the answer starts/ends (counted from top left to bottom right of table)
:param start: Position where the span starts
:param end: Position where the spand ends
:param end: Position where the span ends
"""
def __contains__(self, value):
@@ -334,14 +338,26 @@ class Span:
) from e
@dataclass
class TableCell:
row: int
col: int
"""
Defining a table cell via the row and column index.
:param row: Row index of the cell
:param col: Column index of the cell
"""
@dataclass
class Answer:
answer: str
type: Literal["generative", "extractive", "other"] = "extractive"
score: Optional[float] = None
context: Optional[Union[str, pd.DataFrame]] = None
offsets_in_document: Optional[List[Span]] = None
offsets_in_context: Optional[List[Span]] = None
offsets_in_document: Optional[Union[List[Span], List[TableCell]]] = None
offsets_in_context: Optional[Union[List[Span], List[TableCell]]] = None
document_ids: Optional[List[str]] = None
meta: Optional[Dict[str, Any]] = None
@@ -374,12 +390,13 @@ class Answer:
"""
def __post_init__(self):
# In case offsets are passed as dicts rather than Span objects we convert them here
# In case offsets are passed as dicts rather than Span or TableCell objects we convert them here
# For example, this is used when instantiating an object via from_json()
if self.offsets_in_document is not None:
self.offsets_in_document = [Span(**e) if isinstance(e, dict) else e for e in self.offsets_in_document]
self.offsets_in_document = self._from_dict_offsets(self.offsets_in_document)
if self.offsets_in_context is not None:
self.offsets_in_context = [Span(**e) if isinstance(e, dict) else e for e in self.offsets_in_context]
self.offsets_in_context = self._from_dict_offsets(self.offsets_in_context)
if self.meta is None:
self.meta = {}
@@ -419,6 +436,19 @@ class Answer:
data = json.loads(data)
return cls.from_dict(data)
@staticmethod
def _from_dict_offsets(offsets):
converted_offsets = []
for e in offsets:
if isinstance(e, dict):
if "row" in e: # is a TableCell
converted_offsets.append(TableCell(**e))
else:
converted_offsets.append(Span(**e))
else:
converted_offsets.append(e)
return converted_offsets
@dataclass
class Label:
@@ -623,10 +653,10 @@ class MultiLabel:
for answer in answered:
if answer.offsets_in_document is not None:
for span in answer.offsets_in_document:
self._offsets_in_documents.append({"start": span.start, "end": span.end})
self._offsets_in_documents.append(self._to_dict_offsets(span))
if answer.offsets_in_context is not None:
for span in answer.offsets_in_context:
self._offsets_in_contexts.append({"start": span.start, "end": span.end})
self._offsets_in_contexts.append(self._to_dict_offsets(span))
# There are two options here to represent document_ids:
# taking the id from the document of each label or taking the document_id of each label's answer.
@@ -639,6 +669,13 @@ class MultiLabel:
self._document_ids = [l.document.id for l in self._labels if not l.no_answer]
self._contexts = [str(l.document.content) for l in self._labels if not l.no_answer]
@staticmethod
def _to_dict_offsets(offset: Union[Span, TableCell]) -> Dict:
if isinstance(offset, TableCell):
return {"row": offset.row, "col": offset.col}
else:
return {"start": offset.start, "end": offset.end}
@property
def labels(self):
return self._labels
+16 -1
View File
@@ -13,7 +13,7 @@ import pandas as pd
import pytest
from fastapi.testclient import TestClient
import posthog
from haystack import Document, Answer, Pipeline
from haystack import Document, Answer, Pipeline, TableCell
import haystack
from haystack.nodes import BaseReader, BaseRetriever
from haystack.document_stores import BaseDocumentStore
@@ -435,12 +435,27 @@ def test_query_with_dataframe(client):
meta={"test_key": "test_value"},
)
],
"answers": [
Answer(
answer="text_2",
type="extractive",
score=0.95,
context=pd.DataFrame.from_records([{"col1": "text_1", "col2": 1}, {"col1": "text_2", "col2": 2}]),
offsets_in_document=[TableCell(1, 0)],
offsets_in_context=[TableCell(1, 0)],
meta={"aggregation_operator": "NONE", "answer_cells": ["text_2"]},
)
],
}
response = client.post(url="/query", json={"query": TEST_QUERY})
assert 200 == response.status_code
assert len(response.json()["documents"]) == 1
assert response.json()["documents"][0]["content"] == [["col1", "col2"], ["text_1", 1], ["text_2", 2]]
assert response.json()["documents"][0]["content_type"] == "table"
assert len(response.json()["answers"]) == 1
assert response.json()["answers"][0]["context"] == [["col1", "col2"], ["text_1", 1], ["text_2", 2]]
assert response.json()["answers"][0]["offsets_in_document"] == [{"row": 1, "col": 0}]
assert response.json()["answers"][0]["offsets_in_context"] == [{"row": 1, "col": 0}]
# Ensure `run` was called with the expected parameters
mocked_pipeline.run.assert_called_with(query=TEST_QUERY, params={}, debug=False)
+4 -3
View File
@@ -511,16 +511,17 @@ def reader(request):
@pytest.fixture(params=["tapas_small", "tapas_base", "tapas_scored", "rci"])
def table_reader_and_param(request):
if request.param == "tapas_small":
return TableReader(model_name_or_path="google/tapas-small-finetuned-wtq"), request.param
return TableReader(model_name_or_path="google/tapas-small-finetuned-wtq", return_table_cell=True), request.param
elif request.param == "tapas_base":
return TableReader(model_name_or_path="google/tapas-base-finetuned-wtq"), request.param
return TableReader(model_name_or_path="google/tapas-base-finetuned-wtq", return_table_cell=True), request.param
elif request.param == "tapas_scored":
return TableReader(model_name_or_path="deepset/tapas-large-nq-hn-reader"), request.param
return TableReader(model_name_or_path="deepset/tapas-large-nq-hn-reader", return_table_cell=True), request.param
elif request.param == "rci":
return (
RCIReader(
row_model_name_or_path="michaelrglass/albert-base-rci-wikisql-row",
column_model_name_or_path="michaelrglass/albert-base-rci-wikisql-col",
return_table_cell=True,
),
request.param,
)
+31 -18
View File
@@ -1,12 +1,13 @@
import logging
import pandas as pd
import torch
import pytest
from haystack.schema import Document, Answer
from haystack.schema import Document, Answer, Span, TableCell
from haystack.pipelines.base import Pipeline
from haystack.nodes.reader.table import _calculate_answer_offsets_span, _calculate_answer_offsets
@pytest.fixture
def table_doc1():
@@ -39,6 +40,18 @@ def table_doc3():
return Document(content=pd.DataFrame(data), content_type="table", id="doc3")
@pytest.mark.unit
def test_calculate_answer_offsets_span(table_doc1):
offsets_span = _calculate_answer_offsets_span(answer_coordinates=[(0, 1), (1, 3)], table=table_doc1.content)
assert offsets_span == [Span(start=1, end=2), Span(start=7, end=8)]
@pytest.mark.unit
def test_calculate_answer_offsets_table_cell(table_doc1):
offsets_span = _calculate_answer_offsets(answer_coordinates=[(0, 1), (1, 3)])
assert offsets_span == [TableCell(row=0, col=1), TableCell(row=1, col=3)]
@pytest.mark.integration
@pytest.mark.parametrize("table_reader_and_param", ["tapas_small", "rci", "tapas_scored"], indirect=True)
def test_table_reader(table_reader_and_param, table_doc1, table_doc2):
@@ -57,20 +70,20 @@ def test_table_reader(table_reader_and_param, table_doc1, table_doc2):
reference1 = {"tapas_small": {"score": 1.0}, "rci": {"score": -6.5301}, "tapas_scored": {"score": 0.50568}}
assert prediction["answers"][0].score == pytest.approx(reference1[param]["score"], rel=1e-3)
assert prediction["answers"][0].answer == "11 november 1974"
assert prediction["answers"][0].offsets_in_context[0].start == 7
assert prediction["answers"][0].offsets_in_context[0].end == 8
assert prediction["answers"][0].offsets_in_context[0].row == 1
assert prediction["answers"][0].offsets_in_context[0].col == 3
assert prediction["answers"][0].document_ids == ["doc1"]
# Check the second answer in the list
reference2 = {
"tapas_small": {"answer": "5 april 1980", "start": 7, "end": 8, "score": 0.86314, "doc_id": ["doc2"]},
"rci": {"answer": "47", "start": 5, "end": 6, "score": -6.836, "doc_id": ["doc1"]},
"tapas_scored": {"answer": "brad pitt", "start": 0, "end": 1, "score": 0.49078, "doc_id": ["doc1"]},
"tapas_small": {"answer": "5 april 1980", "row": 1, "col": 3, "score": 0.86314, "doc_id": ["doc2"]},
"rci": {"answer": "47", "row": 1, "col": 1, "score": -6.836, "doc_id": ["doc1"]},
"tapas_scored": {"answer": "brad pitt", "row": 0, "col": 0, "score": 0.49078, "doc_id": ["doc1"]},
}
assert prediction["answers"][1].score == pytest.approx(reference2[param]["score"], rel=1e-3)
assert prediction["answers"][1].answer == reference2[param]["answer"]
assert prediction["answers"][1].offsets_in_context[0].start == reference2[param]["start"]
assert prediction["answers"][1].offsets_in_context[0].end == reference2[param]["end"]
assert prediction["answers"][1].offsets_in_context[0].row == reference2[param]["row"]
assert prediction["answers"][1].offsets_in_context[0].col == reference2[param]["col"]
assert prediction["answers"][1].document_ids == reference2[param]["doc_id"]
@@ -100,20 +113,20 @@ def test_table_reader_batch_single_query_single_doc_list(table_reader_and_param,
score_reference = {"tapas_small": {"score": 1.0}, "rci": {"score": -6.5301}, "tapas_scored": {"score": 0.50568}}
assert prediction["answers"][0][0].score == pytest.approx(score_reference[param]["score"], rel=1e-3)
assert prediction["answers"][0][0].answer == "11 november 1974"
assert prediction["answers"][0][0].offsets_in_context[0].start == 7
assert prediction["answers"][0][0].offsets_in_context[0].end == 8
assert prediction["answers"][0][0].offsets_in_context[0].row == 1
assert prediction["answers"][0][0].offsets_in_context[0].col == 3
assert prediction["answers"][0][0].document_ids == ["doc1"]
# Check first answer from the 2ND Document
ans_reference = {
"tapas_small": {"answer": "5 april 1980", "start": 7, "end": 8, "score": 0.86314, "doc_id": ["doc2"]},
"rci": {"answer": "15 september 1960", "start": 11, "end": 12, "score": -7.9429, "doc_id": ["doc2"]},
"tapas_scored": {"answer": "5", "start": 10, "end": 11, "score": 0.11485, "doc_id": ["doc2"]},
"tapas_small": {"answer": "5 april 1980", "row": 1, "col": 3, "score": 0.86314, "doc_id": ["doc2"]},
"rci": {"answer": "15 september 1960", "row": 2, "col": 3, "score": -7.9429, "doc_id": ["doc2"]},
"tapas_scored": {"answer": "5", "row": 2, "col": 2, "score": 0.11485, "doc_id": ["doc2"]},
}
assert prediction["answers"][1][0].score == pytest.approx(ans_reference[param]["score"], rel=1e-3)
assert prediction["answers"][1][0].answer == ans_reference[param]["answer"]
assert prediction["answers"][1][0].offsets_in_context[0].start == ans_reference[param]["start"]
assert prediction["answers"][1][0].offsets_in_context[0].end == ans_reference[param]["end"]
assert prediction["answers"][1][0].offsets_in_context[0].row == ans_reference[param]["row"]
assert prediction["answers"][1][0].offsets_in_context[0].col == ans_reference[param]["col"]
assert prediction["answers"][1][0].document_ids == ans_reference[param]["doc_id"]
@@ -207,8 +220,8 @@ def test_table_reader_in_pipeline(table_reader_and_param, table_doc1):
prediction = pipeline.run(query=query, documents=[table_doc1])
assert prediction["answers"][0].answer == "11 november 1974"
assert prediction["answers"][0].offsets_in_context[0].start == 7
assert prediction["answers"][0].offsets_in_context[0].end == 8
assert prediction["answers"][0].offsets_in_context[0].row == 1
assert prediction["answers"][0].offsets_in_context[0].col == 3
assert prediction["answers"][0].document_ids == ["doc1"]
+172 -90
View File
@@ -1,53 +1,86 @@
from haystack.schema import Document, Label, Answer, Span, MultiLabel
from haystack.schema import Document, Label, Answer, Span, MultiLabel, TableCell
import pytest
import numpy as np
import pandas as pd
LABELS = [
Label(
query="some",
answer=Answer(
answer="an answer",
type="extractive",
score=0.1,
document_ids=["123"],
offsets_in_document=[Span(start=1, end=3)],
@pytest.fixture
def text_labels():
return [
Label(
query="some",
answer=Answer(
answer="an answer",
type="extractive",
score=0.1,
document_ids=["123"],
offsets_in_document=[Span(start=1, end=3)],
),
document=Document(content="some text", content_type="text"),
is_correct_answer=True,
is_correct_document=True,
origin="user-feedback",
),
document=Document(content="some text", content_type="text"),
is_correct_answer=True,
is_correct_document=True,
origin="user-feedback",
),
Label(
query="some",
answer=Answer(answer="annother answer", type="extractive", score=0.1, document_ids=["123"]),
document=Document(content="some text", content_type="text"),
is_correct_answer=True,
is_correct_document=True,
origin="user-feedback",
),
Label(
query="some",
answer=Answer(
answer="an answer",
type="extractive",
score=0.1,
document_ids=["123"],
offsets_in_document=[Span(start=1, end=3)],
Label(
query="some",
answer=Answer(answer="annother answer", type="extractive", score=0.1, document_ids=["123"]),
document=Document(content="some text", content_type="text"),
is_correct_answer=True,
is_correct_document=True,
origin="user-feedback",
),
document=Document(content="some text", content_type="text"),
is_correct_answer=True,
is_correct_document=True,
origin="user-feedback",
),
]
Label(
query="some",
answer=Answer(
answer="an answer",
type="extractive",
score=0.1,
document_ids=["123"],
offsets_in_document=[Span(start=1, end=3)],
),
document=Document(content="some text", content_type="text"),
is_correct_answer=True,
is_correct_document=True,
origin="user-feedback",
),
]
def test_document_from_dict():
doc = Document(
content="this is the content of the document", meta={"some": "meta"}, id_hash_keys=["content", "meta"]
@pytest.fixture
def text_answer():
return Answer(
answer="an answer",
type="extractive",
score=0.1,
context="abc",
offsets_in_document=[Span(start=1, end=10)],
offsets_in_context=[Span(start=3, end=5)],
document_ids=["123"],
)
@pytest.fixture
def table_doc():
data = {
"actors": ["brad pitt", "leonardo di caprio", "george clooney"],
"age": [58, 47, 60],
"number of movies": [87, 53, 69],
"date of birth": ["18 december 1963", "11 november 1974", "6 may 1961"],
}
return Document(content=pd.DataFrame(data), content_type="table", id="doc1")
@pytest.fixture
def table_doc_with_embedding():
data = {
"actors": ["brad pitt", "leonardo di caprio", "george clooney"],
"age": [58, 47, 60],
"number of movies": [87, 53, 69],
"date of birth": ["18 december 1963", "11 november 1974", "6 may 1961"],
}
return Document(
content=pd.DataFrame(data), content_type="table", id="doc2", embedding=np.random.rand(768).astype(np.float32)
)
assert doc == Document.from_dict(doc.to_dict())
def test_no_answer_label():
@@ -92,21 +125,27 @@ def test_no_answer_label():
assert labels[3].no_answer == False
def test_equal_label():
assert LABELS[2] == LABELS[0]
assert LABELS[1] != LABELS[0]
def test_equal_label(text_labels):
assert text_labels[2] == text_labels[0]
assert text_labels[1] != text_labels[0]
def test_answer_to_json():
a = Answer(
answer="an answer",
type="extractive",
score=0.1,
context="abc",
offsets_in_document=[Span(start=1, end=10)],
offsets_in_context=[Span(start=3, end=5)],
document_ids=["123"],
)
def test_label_to_json(text_labels):
j0 = text_labels[0].to_json()
l_new = Label.from_json(j0)
assert l_new == text_labels[0]
assert l_new.answer.offsets_in_document[0].start == 1
def test_label_to_dict(text_labels):
j0 = text_labels[0].to_dict()
l_new = Label.from_dict(j0)
assert l_new == text_labels[0]
assert l_new.answer.offsets_in_document[0].start == 1
def test_answer_to_json(text_answer):
a = text_answer
j = a.to_json()
assert type(j) == str
assert len(j) > 30
@@ -115,16 +154,8 @@ def test_answer_to_json():
assert a_new == a
def test_answer_to_dict():
a = Answer(
answer="an answer",
type="extractive",
score=0.1,
context="abc",
offsets_in_document=[Span(start=1, end=10)],
offsets_in_context=[Span(start=3, end=5)],
document_ids=["123"],
)
def test_answer_to_dict(text_answer):
a = text_answer
j = a.to_dict()
assert type(j) == dict
a_new = Answer.from_dict(j)
@@ -132,24 +163,15 @@ def test_answer_to_dict():
assert a_new == a
def test_label_to_json():
j0 = LABELS[0].to_json()
l_new = Label.from_json(j0)
assert l_new == LABELS[0]
def test_document_from_dict():
doc = Document(
content="this is the content of the document", meta={"some": "meta"}, id_hash_keys=["content", "meta"]
)
assert doc == Document.from_dict(doc.to_dict())
def test_label_to_json():
j0 = LABELS[0].to_json()
l_new = Label.from_json(j0)
assert l_new == LABELS[0]
assert l_new.answer.offsets_in_document[0].start == 1
def test_label_to_dict():
j0 = LABELS[0].to_dict()
l_new = Label.from_dict(j0)
assert l_new == LABELS[0]
assert l_new.answer.offsets_in_document[0].start == 1
def test_table_document_from_dict(table_doc):
assert table_doc == Document.from_dict(table_doc.to_dict())
def test_doc_to_json():
@@ -180,12 +202,30 @@ def test_doc_to_json():
assert d == d_new
def test_table_doc_to_json(table_doc, table_doc_with_embedding):
# With embedding
j0 = table_doc_with_embedding.to_json()
d_new = Document.from_json(j0)
assert table_doc_with_embedding == d_new
# No embedding
j0 = table_doc.to_json()
d_new = Document.from_json(j0)
assert table_doc == d_new
def test_answer_postinit():
a = Answer(answer="test", offsets_in_document=[{"start": 10, "end": 20}])
assert a.meta == {}
assert isinstance(a.offsets_in_document[0], Span)
def test_table_answer_postinit():
a = Answer(answer="test", offsets_in_document=[{"row": 1, "col": 2}])
assert a.meta == {}
assert isinstance(a.offsets_in_document[0], TableCell)
def test_generate_doc_id_using_text():
text1 = "text1"
text2 = "text2"
@@ -334,7 +374,11 @@ def test_multilabel_preserve_order_w_duplicates():
Label(
id="0",
query="question",
answer=Answer(answer="answer1", offsets_in_document=[Span(start=12, end=18)]),
answer=Answer(
answer="answer1",
offsets_in_document=[Span(start=12, end=18)],
offsets_in_context=[Span(start=1, end=7)],
),
document=Document(content="some", id="123"),
is_correct_answer=True,
is_correct_document=True,
@@ -343,7 +387,11 @@ def test_multilabel_preserve_order_w_duplicates():
Label(
id="1",
query="question",
answer=Answer(answer="answer2", offsets_in_document=[Span(start=12, end=18)]),
answer=Answer(
answer="answer2",
offsets_in_document=[Span(start=10, end=16)],
offsets_in_context=[Span(start=0, end=6)],
),
document=Document(content="some", id="123"),
is_correct_answer=True,
is_correct_document=True,
@@ -352,7 +400,11 @@ def test_multilabel_preserve_order_w_duplicates():
Label(
id="2",
query="question",
answer=Answer(answer="answer3", offsets_in_document=[Span(start=12, end=18)]),
answer=Answer(
answer="answer3",
offsets_in_document=[Span(start=14, end=20)],
offsets_in_context=[Span(start=2, end=8)],
),
document=Document(content="some other", id="333"),
is_correct_answer=True,
is_correct_document=True,
@@ -361,7 +413,11 @@ def test_multilabel_preserve_order_w_duplicates():
Label(
id="0",
query="question",
answer=Answer(answer="answer1", offsets_in_document=[Span(start=12, end=18)]),
answer=Answer(
answer="answer1",
offsets_in_document=[Span(start=12, end=18)],
offsets_in_context=[Span(start=1, end=7)],
),
document=Document(content="some", id="123"),
is_correct_answer=True,
is_correct_document=True,
@@ -370,7 +426,11 @@ def test_multilabel_preserve_order_w_duplicates():
Label(
id="2",
query="question",
answer=Answer(answer="answer3", offsets_in_document=[Span(start=12, end=18)]),
answer=Answer(
answer="answer3",
offsets_in_document=[Span(start=14, end=20)],
offsets_in_context=[Span(start=2, end=8)],
),
document=Document(content="some other", id="333"),
is_correct_answer=True,
is_correct_document=True,
@@ -380,7 +440,16 @@ def test_multilabel_preserve_order_w_duplicates():
multilabel = MultiLabel(labels=labels)
assert len(multilabel.document_ids) == 3
assert multilabel.query == "question"
assert multilabel.answers == ["answer1", "answer2", "answer3"]
assert multilabel.document_ids == ["123", "123", "333"]
assert multilabel.contexts == ["some", "some", "some other"]
assert multilabel.offsets_in_documents == [
{"start": 12, "end": 18},
{"start": 10, "end": 16},
{"start": 14, "end": 20},
]
assert multilabel.offsets_in_contexts == [{"start": 1, "end": 7}, {"start": 0, "end": 6}, {"start": 2, "end": 8}]
for i in range(0, 3):
assert multilabel.labels[i].id == str(i)
@@ -427,16 +496,29 @@ def test_multilabel_id():
def test_multilabel_with_doc_containing_dataframes():
table = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
table_doc = Document(content=table, content_type="table", id="table1")
label = Label(
query="A question",
document=Document(content=pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})),
document=table_doc,
is_correct_answer=True,
is_correct_document=True,
origin="gold-label",
answer=Answer(answer="answer 1"),
answer=Answer(
answer="1",
context=table,
offsets_in_document=[TableCell(0, 0)],
offsets_in_context=[TableCell(0, 0)],
document_ids=[table_doc.id],
),
)
assert len(MultiLabel(labels=[label]).contexts) == 1
assert type(MultiLabel(labels=[label]).contexts[0]) is str
multilabel = MultiLabel(labels=[label])
assert multilabel.query == "A question"
assert multilabel.contexts == [" col1 col2\n0 1 3\n1 2 4"]
assert multilabel.answers == ["1"]
assert multilabel.document_ids == ["table1"]
assert multilabel.offsets_in_documents == [{"row": 0, "col": 0}]
assert multilabel.offsets_in_contexts == [{"row": 0, "col": 0}]
def test_multilabel_serialization():
+4 -4
View File
@@ -25,7 +25,7 @@ from haystack.pipelines.standard_pipelines import (
TranslationWrapperPipeline,
)
from haystack.nodes.translator.transformers import TransformersTranslator
from haystack.schema import Answer, Document, EvaluationResult, Label, MultiLabel, Span
from haystack.schema import Answer, Document, EvaluationResult, Label, MultiLabel, Span, TableCell
@pytest.mark.skipif(sys.platform in ["win32", "cygwin"], reason="Causes OOM on windows github runner")
@@ -415,7 +415,7 @@ EVAL_TABLE_LABELS = [
labels=[
Label(
query="How old is Brad Pitt?",
answer=Answer(answer="56", offsets_in_context=[Span(1, 2)]),
answer=Answer(answer="56", offsets_in_context=[TableCell(1, 2)]),
document=Document(
id="a044cf3fb8aade03a12399c7a2fe9a6b",
content_type="table",
@@ -434,7 +434,7 @@ EVAL_TABLE_LABELS = [
),
Label( # Label with different doc but same answer and query
query="How old is Brad Pitt?",
answer=Answer(answer="56", offsets_in_context=[Span(4, 5)]),
answer=Answer(answer="56", offsets_in_context=[TableCell(4, 5)]),
document=Document(
id="a044cf3fb8aade03a12399c7a2fe9a6b",
content_type="table",
@@ -453,7 +453,7 @@ EVAL_TABLE_LABELS = [
labels=[
Label(
query="To which state does Spikeroog belong?",
answer=Answer(answer="Lower Saxony", offsets_in_context=[Span(7, 8)]),
answer=Answer(answer="Lower Saxony", offsets_in_context=[TableCell(7, 8)]),
document=Document(
id="b044cf3fb8aade03a12399c7a2fe9a6c",
content_type="table",