|
import json |
|
import os |
|
from dataclasses import asdict, dataclass, is_dataclass |
|
from itertools import chain |
|
from typing import Dict, FrozenSet, List, Set, Tuple, Union |
|
|
|
|
|
@dataclass(eq=True, frozen=True) |
|
class Evidence: |
|
""" |
|
(docid, start_token, end_token) form the only official Evidence; sentence level annotations are for convenience. |
|
Args: |
|
text: Some representation of the evidence text |
|
docid: Some identifier for the document |
|
start_token: The canonical start token, inclusive |
|
end_token: The canonical end token, exclusive |
|
start_sentence: Best guess start sentence, inclusive |
|
end_sentence: Best guess end sentence, exclusive |
|
""" |
|
|
|
text: Union[str, Tuple[int], Tuple[str]] |
|
docid: str |
|
start_token: int = -1 |
|
end_token: int = -1 |
|
start_sentence: int = -1 |
|
end_sentence: int = -1 |
|
|
|
|
|
@dataclass(eq=True, frozen=True) |
|
class Annotation: |
|
""" |
|
Args: |
|
annotation_id: unique ID for this annotation element |
|
query: some representation of a query string |
|
evidences: a set of "evidence groups". |
|
Each evidence group is: |
|
* sufficient to respond to the query (or justify an answer) |
|
* composed of one or more Evidences |
|
* may have multiple documents in it (depending on the dataset) |
|
- e-snli has multiple documents |
|
- other datasets do not |
|
classification: str |
|
query_type: Optional str, additional information about the query |
|
docids: a set of docids in which one may find evidence. |
|
""" |
|
|
|
annotation_id: str |
|
query: Union[str, Tuple[int]] |
|
evidences: Union[Set[Tuple[Evidence]], FrozenSet[Tuple[Evidence]]] |
|
classification: str |
|
query_type: str = None |
|
docids: Set[str] = None |
|
|
|
def all_evidences(self) -> Tuple[Evidence]: |
|
return tuple(list(chain.from_iterable(self.evidences))) |
|
|
|
|
|
def annotations_to_jsonl(annotations, output_file): |
|
with open(output_file, "w") as of: |
|
for ann in sorted(annotations, key=lambda x: x.annotation_id): |
|
as_json = _annotation_to_dict(ann) |
|
as_str = json.dumps(as_json, sort_keys=True) |
|
of.write(as_str) |
|
of.write("\n") |
|
|
|
|
|
def _annotation_to_dict(dc): |
|
|
|
if is_dataclass(dc): |
|
d = asdict(dc) |
|
ret = dict() |
|
for k, v in d.items(): |
|
ret[k] = _annotation_to_dict(v) |
|
return ret |
|
elif isinstance(dc, dict): |
|
ret = dict() |
|
for k, v in dc.items(): |
|
k = _annotation_to_dict(k) |
|
v = _annotation_to_dict(v) |
|
ret[k] = v |
|
return ret |
|
elif isinstance(dc, str): |
|
return dc |
|
elif isinstance(dc, (set, frozenset, list, tuple)): |
|
ret = [] |
|
for x in dc: |
|
ret.append(_annotation_to_dict(x)) |
|
return tuple(ret) |
|
else: |
|
return dc |
|
|
|
|
|
def load_jsonl(fp: str) -> List[dict]: |
|
ret = [] |
|
with open(fp, "r") as inf: |
|
for line in inf: |
|
content = json.loads(line) |
|
ret.append(content) |
|
return ret |
|
|
|
|
|
def write_jsonl(jsonl, output_file): |
|
with open(output_file, "w") as of: |
|
for js in jsonl: |
|
as_str = json.dumps(js, sort_keys=True) |
|
of.write(as_str) |
|
of.write("\n") |
|
|
|
|
|
def annotations_from_jsonl(fp: str) -> List[Annotation]: |
|
ret = [] |
|
with open(fp, "r") as inf: |
|
for line in inf: |
|
content = json.loads(line) |
|
ev_groups = [] |
|
for ev_group in content["evidences"]: |
|
ev_group = tuple([Evidence(**ev) for ev in ev_group]) |
|
ev_groups.append(ev_group) |
|
content["evidences"] = frozenset(ev_groups) |
|
ret.append(Annotation(**content)) |
|
return ret |
|
|
|
|
|
def load_datasets( |
|
data_dir: str, |
|
) -> Tuple[List[Annotation], List[Annotation], List[Annotation]]: |
|
"""Loads a training, validation, and test dataset |
|
|
|
Each dataset is assumed to have been serialized by annotations_to_jsonl, |
|
that is it is a list of json-serialized Annotation instances. |
|
""" |
|
train_data = annotations_from_jsonl(os.path.join(data_dir, "train.jsonl")) |
|
val_data = annotations_from_jsonl(os.path.join(data_dir, "val.jsonl")) |
|
test_data = annotations_from_jsonl(os.path.join(data_dir, "test.jsonl")) |
|
return train_data, val_data, test_data |
|
|
|
|
|
def load_documents( |
|
data_dir: str, docids: Set[str] = None |
|
) -> Dict[str, List[List[str]]]: |
|
"""Loads a subset of available documents from disk. |
|
|
|
Each document is assumed to be serialized as newline ('\n') separated sentences. |
|
Each sentence is assumed to be space (' ') joined tokens. |
|
""" |
|
if os.path.exists(os.path.join(data_dir, "docs.jsonl")): |
|
assert not os.path.exists(os.path.join(data_dir, "docs")) |
|
return load_documents_from_file(data_dir, docids) |
|
|
|
docs_dir = os.path.join(data_dir, "docs") |
|
res = dict() |
|
if docids is None: |
|
docids = sorted(os.listdir(docs_dir)) |
|
else: |
|
docids = sorted(set(str(d) for d in docids)) |
|
for d in docids: |
|
with open(os.path.join(docs_dir, d), "r") as inf: |
|
res[d] = inf.read() |
|
return res |
|
|
|
|
|
def load_flattened_documents(data_dir: str, docids: Set[str]) -> Dict[str, List[str]]: |
|
"""Loads a subset of available documents from disk. |
|
|
|
Returns a tokenized version of the document. |
|
""" |
|
unflattened_docs = load_documents(data_dir, docids) |
|
flattened_docs = dict() |
|
for doc, unflattened in unflattened_docs.items(): |
|
flattened_docs[doc] = list(chain.from_iterable(unflattened)) |
|
return flattened_docs |
|
|
|
|
|
def intern_documents( |
|
documents: Dict[str, List[List[str]]], word_interner: Dict[str, int], unk_token: str |
|
): |
|
""" |
|
Replaces every word with its index in an embeddings file. |
|
|
|
If a word is not found, uses the unk_token instead |
|
""" |
|
ret = dict() |
|
unk = word_interner[unk_token] |
|
for docid, sentences in documents.items(): |
|
ret[docid] = [[word_interner.get(w, unk) for w in s] for s in sentences] |
|
return ret |
|
|
|
|
|
def intern_annotations( |
|
annotations: List[Annotation], word_interner: Dict[str, int], unk_token: str |
|
): |
|
ret = [] |
|
for ann in annotations: |
|
ev_groups = [] |
|
for ev_group in ann.evidences: |
|
evs = [] |
|
for ev in ev_group: |
|
evs.append( |
|
Evidence( |
|
text=tuple( |
|
[ |
|
word_interner.get(t, word_interner[unk_token]) |
|
for t in ev.text.split() |
|
] |
|
), |
|
docid=ev.docid, |
|
start_token=ev.start_token, |
|
end_token=ev.end_token, |
|
start_sentence=ev.start_sentence, |
|
end_sentence=ev.end_sentence, |
|
) |
|
) |
|
ev_groups.append(tuple(evs)) |
|
ret.append( |
|
Annotation( |
|
annotation_id=ann.annotation_id, |
|
query=tuple( |
|
[ |
|
word_interner.get(t, word_interner[unk_token]) |
|
for t in ann.query.split() |
|
] |
|
), |
|
evidences=frozenset(ev_groups), |
|
classification=ann.classification, |
|
query_type=ann.query_type, |
|
) |
|
) |
|
return ret |
|
|
|
|
|
def load_documents_from_file( |
|
data_dir: str, docids: Set[str] = None |
|
) -> Dict[str, List[List[str]]]: |
|
"""Loads a subset of available documents from 'docs.jsonl' file on disk. |
|
|
|
Each document is assumed to be serialized as newline ('\n') separated sentences. |
|
Each sentence is assumed to be space (' ') joined tokens. |
|
""" |
|
docs_file = os.path.join(data_dir, "docs.jsonl") |
|
documents = load_jsonl(docs_file) |
|
documents = {doc["docid"]: doc["document"] for doc in documents} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
return documents |
|
|
