epsilon-compositionality

extract_similarities.py (5691B)

---

 from typing import Any, NamedTuple, Iterator
 import copy
 import json
 import pathlib
 
 import torch
 import tqdm
 
 from epsilon_compositionality.models import ModelFactory, Model
 
 
 INPUT_PATH = pathlib.Path(__file__).parent / "data" / "ncimp.processed.jsonl"
 OUTPUT_PATH = pathlib.Path(__file__).parent / "data" / "similarities"
 
 
 def process_sample(input_sample: dict[str, Any]) -> dict[str, Any] | None:
     """Process a NCIMP sample into set of (a, b, ac, bc).
 
     A processed sample has the following shape:
     {
         "compound": "acid test",
         "comp_class": "NC",
         "head": [
             { # sentence 1
                 "positives": [ # compare "acid test" to alternatives
                     {
                         "phrase_a": "acid test",
                         "phrase_b": "acid trial",
                         "context_ac": "sentence1 with acid test…",
                         "context_bc": "sentence1 with acid trial…"
                     },
                     {
                         "phrase_a": "acid test",
                         "phrase_b": "acid analysis",
                         "context_ac": "sentence1 with acid test…",
                         "context_bc": "sentence1 with acid analysis…"
                     },
                     # …
                 ],
                 "negatives": [ # compare alternatives between themselves
                     {
                         "phrase_a": "acid analysis",
                         "phrase_b": "acid trial",
                         "context_ac": "sentence1 with acid analysis…",
                         "context_bc": "sentence1 with acid trial…"
                     },
                     {
                         "phrase_a": "acid tryout",
                         "phrase_b": "acid analysis",
                         "context_ac": "sentence1 with acid tryout…",
                         "context_bc": "sentence1 with acid analysis…"
                     },
                 ]
             },
             # … more sentences
         ],
         "modifier": [
             # similarly set of sentences where acid is replaced instead of test (e.g. sour test)
         ]
     }
     """
     output_sample: dict[str, Any] = {
             "compound": input_sample["compound"],
             "comp_class": input_sample["general class"],
         }
     for side_index, side in enumerate(["modifier", "head"]):
         output_sample[side] = []
         for sentence in input_sample["sentences"]:
             if len(sentence[f"{side} synonym"]) <= 1:
                 continue
 
             alternatives = list(sentence[f"{side} synonym"].items())
             second_alternatives = alternatives[1:] + [alternatives[0]]
             if len(alternatives) == 2:
                 alternatives, second_alternatives = [alternatives[0]], [alternatives[1]]
 
             positives: list[Any] = []
             negatives: list[Any] = []
             for (alt_phrase, alt_sentence), (second_alt_phrase, second_alt_sentence) in zip(alternatives, second_alternatives):
                 positives.append({
                         "phrase_a": input_sample["compound"].split()[side_index],
                         "phrase_b": alt_phrase.split()[side_index],
                         "context_ac": sentence["original"],
                         "context_bc": alt_sentence,
                     })
                 negatives.append({
                         "phrase_a": second_alt_phrase.split()[side_index],
                         "phrase_b": alt_phrase.split()[side_index],
                         "context_ac": second_alt_sentence,
                         "context_bc": alt_sentence,
                     })
             output_sample[side].append({"positives": positives, "negatives": negatives})
         if not output_sample[side]:
             return None
     return output_sample
 
 
 def load_dataset() -> list[dict[str, Any]]:
     dataset: list[dict[str, Any]] = []
     with INPUT_PATH.open("r") as input_file:
         for line in input_file:
             sample: dict[str, Any] | None = process_sample(json.loads(line))
             if sample:
                 dataset.append(sample)
     return dataset
 
 
 def main() -> None:
     dataset: list[dict[str, Any]] = load_dataset()
     OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
     for model_name in ModelFactory._factory:
         model_path: pathlib.Path = OUTPUT_PATH / f"{model_name}.jsonl"
 
         model: Model = ModelFactory.get(model_name)
         output: list[dict[str, Any]] = []
 
         with model_path.open("w") as output_file:
             for input_sample in tqdm.tqdm(dataset, f"Embedding dataset with {model_name}"):
                 output_sample: dict[str, Any] = copy.deepcopy(input_sample)
                 for side_index, side in enumerate(["modifier", "head"]):
                     for sentence in output_sample[side]:
                         for polarity in ["positives", "negatives"]:
                             for subsample in sentence[polarity]:
                                 phrase_a: torch.Tensor = model.embed(subsample["phrase_a"])
                                 phrase_b: torch.Tensor = model.embed(subsample["phrase_b"])
                                 context_ac: torch.Tensor = model.embed(subsample["context_ac"])
                                 context_bc: torch.Tensor = model.embed(subsample["context_bc"])
                                 subsample["phrase_similarity"] = model.calculate_similarity(phrase_a, phrase_b)
                                 subsample["contextual_similarity"] = model.calculate_similarity(context_ac, context_bc)
                 print(json.dumps(output_sample), file=output_file)
 
 
 if __name__ == "__main__":
     main()