build_dataset.py (4708B)
1 #!/usr/bin/env python 2 3 from typing import Any 4 import json 5 import pathlib 6 7 import pandas as pd 8 9 10 INPUT_PATH = pathlib.Path(__file__).parent / "data" / "Finding-Idiomaticity-in-Word-Representations" / "dataset" 11 OUTPUT_PATH = pathlib.Path(__file__).parent / "data" / "ncimp.processed.jsonl" 12 13 14 def sample_filter(frame: pd.DataFrame) -> pd.DataFrame: 15 return frame[(frame["language"] == "en") & (frame["experiment_type"] == "Naturalistic")] 16 17 18 def extract_tagged(sentence: str, tag: str) -> str: 19 return " ".join(word for word, annotation in zip(sentence.split(), eval(tag)) if annotation) 20 21 22 def fix_data(compound: str, keep_modifier: True, keep_head: True, alternatives: dict[str, str]) -> dict[str, str]: 23 fixed: dict[str, str] = {} 24 25 process = lambda x: " ".join(x.lower().split()[1 - keep_modifier:1 + keep_head]) 26 27 uniques: set[str] = { process(compound), f"{process(compound)}s" } 28 for key, value in alternatives.items(): 29 if process(key) not in uniques and f"{process(key)}s" not in uniques: 30 uniques.add(process(key)) 31 uniques.add(f"{process(key)}s") 32 fixed[key] = value 33 34 return fixed 35 36 37 def resolve_sentence(compound: str, input_sentence: pd.Series, support_annotations: dict[str, pd.DataFrame]) -> tuple[dict[str, Any], dict[str, Any]]: 38 filename: str = pathlib.Path(input_sentence["file"]).name 39 file_id: int = int(filename[len("naturalistics_examplesent"):-len(".csv")]) 40 support_annotation: pd.Series = support_annotations[filename].loc[compound] 41 42 common: dict[str, Any] = { 43 "compound": compound, 44 "general class": input_sentence["ClassType"], 45 "general human compositionality": input_sentence["CompositionalityType"], 46 } 47 48 sentence: dict[str, Any] = { 49 "compound": extract_tagged(support_annotation["original sentence"], support_annotation["original sentence_tag"]), 50 "modifier": support_annotation[f"modifier_{file_id}"], # acid 51 "head": support_annotation[f"head_{file_id}"], # test 52 "contextual class": input_sentence["ClassToken"], 53 "contextual human compositionality": input_sentence["CompositionalityToken"], 54 "original": support_annotation["original sentence"], # acid test 55 "modifier synonym": fix_data(compound, True, False, { # sour test 56 extract_tagged(support_annotation[key], support_annotation[f"{key}_tag"]): support_annotation[key] 57 for key in [ 58 f"synonym modifier_{file_id}", 59 f"modifier synonym alt1_{file_id}", 60 f"modifier synonym alt2_{file_id}", 61 f"modifier synonym alt3_{file_id}", 62 f"modifier synonym alt4_{file_id}", 63 ] 64 }), 65 "head synonym": fix_data(compound, False, True, { # acid analysis 66 extract_tagged(support_annotation[key], support_annotation[f"{key}_tag"]): support_annotation[key] 67 for key in [ 68 f"synonym head_{file_id}", 69 f"head synonym alt1_{file_id}", 70 f"head synonym alt2_{file_id}", 71 f"head synonym alt3_{file_id}", 72 f"head synonym alt4_{file_id}", 73 ] 74 }), 75 "both synonym": fix_data(compound, True, True, { # sour analysis 76 extract_tagged(support_annotation[key], support_annotation[f"{key}_tag"]): support_annotation[key] 77 for key in [ 78 f"synonym both", 79 f"both synonym alt1_{file_id}", 80 f"both synonym alt2_{file_id}", 81 f"both synonym alt3_{file_id}", 82 f"both synonym alt4_{file_id}", 83 ] 84 }), 85 } 86 87 return common, sentence 88 89 90 def main() -> None: 91 input_annotations: pd.DataFrame = pd.read_excel(INPUT_PATH / "human_compositionality scores.xlsx") 92 input_annotations = sample_filter(input_annotations) 93 support_annotations: dict[str, pd.DataFrame] = { 94 path.name: pd.read_csv(path).set_index("compound") 95 for path in (INPUT_PATH / "EN").iterdir() 96 if path.name.startswith("naturalistics_examplesent") 97 } 98 99 with OUTPUT_PATH.open("w") as output_file: 100 for compound, input_sample in input_annotations.groupby("compound"): 101 sentences: list[dict[str, Any]] = [] 102 sample: dict[str, Any] = {"sentences": sentences} 103 for input_sentence in input_sample.iloc: 104 sample_data, sentence = resolve_sentence(compound, input_sentence, support_annotations) 105 sentences.append(sentence) 106 sample.update(sample_data) 107 print(json.dumps(sample), file=output_file) 108 109 110 if __name__ == "__main__": 111 main()