diff --git a/haystack/__init__.py b/haystack/__init__.py index 9eeaa7db7b..5da725e7e0 100644 --- a/haystack/__init__.py +++ b/haystack/__init__.py @@ -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 diff --git a/haystack/nodes/reader/farm.py b/haystack/nodes/reader/farm.py index 55ffb0eed6..9e9d065a2d 100644 --- a/haystack/nodes/reader/farm.py +++ b/haystack/nodes/reader/farm.py @@ -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, diff --git a/haystack/nodes/reader/table.py b/haystack/nodes/reader/table.py index 1cd31b2383..85e6674bb2 100644 --- a/haystack/nodes/reader/table.py +++ b/haystack/nodes/reader/table.py @@ -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. diff --git a/haystack/pipelines/base.py b/haystack/pipelines/base.py index 41b1bea145..063f3c5804 100644 --- a/haystack/pipelines/base.py +++ b/haystack/pipelines/base.py @@ -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, diff --git a/haystack/schema.py b/haystack/schema.py index 73ad83168f..46b198021c 100644 --- a/haystack/schema.py +++ b/haystack/schema.py @@ -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 diff --git a/rest_api/test/test_rest_api.py b/rest_api/test/test_rest_api.py index 1b4c0cd786..e0e7a3b833 100644 --- a/rest_api/test/test_rest_api.py +++ b/rest_api/test/test_rest_api.py @@ -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) diff --git a/test/conftest.py b/test/conftest.py index 397c2d0ebd..db341028b4 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -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, ) diff --git a/test/nodes/test_table_reader.py b/test/nodes/test_table_reader.py index 4029806c02..609207ebf6 100644 --- a/test/nodes/test_table_reader.py +++ b/test/nodes/test_table_reader.py @@ -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"] diff --git a/test/others/test_schema.py b/test/others/test_schema.py index df1684f176..dd1553a2f6 100644 --- a/test/others/test_schema.py +++ b/test/others/test_schema.py @@ -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(): diff --git a/test/pipelines/test_eval.py b/test/pipelines/test_eval.py index b57b09a3be..a9a8c61947 100644 --- a/test/pipelines/test_eval.py +++ b/test/pipelines/test_eval.py @@ -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",