| from typing import Dict, List, Tuple, Set
|
| import argparse
|
| import collections
|
| import itertools
|
| import re
|
| from tqdm import tqdm
|
|
|
|
|
| def index_vocabulary(vocabulary_path: str) -> Tuple[Dict[str, int], Dict[int, str]]:
|
| word2index, index2word = dict(), dict()
|
| with open(vocabulary_path) as f:
|
| for i, line in enumerate(f):
|
| word = line.strip()
|
| word2index[word] = i
|
| index2word[i] = word
|
| return word2index, index2word
|
|
|
|
|
| def index_corpus(word2index: Dict[str, int], corpus_path: str) -> Dict[int, Dict[int, List[int]]]:
|
| word_corpus_index = collections.defaultdict(dict)
|
| with open(corpus_path) as f:
|
| for doc_id, line in enumerate(f):
|
| for word_position, word in enumerate(line.strip().split(" ")):
|
| if word in word2index:
|
| wi = word2index[word]
|
| if wi not in word_corpus_index:
|
| word_corpus_index[wi] = collections.defaultdict(list)
|
| word_corpus_index[wi][doc_id].append(word_position)
|
| return word_corpus_index
|
|
|
|
|
| def compute_corpus_occurrences(
|
| word2index: Dict[str, int], word_corpus_index: Dict[int, Dict[int, List[int]]]
|
| ) -> Dict[int, Set[int]]:
|
| corpus_occurrences = collections.defaultdict(set)
|
| mws_splitter = re.compile("[_ ]")
|
|
|
| for word, index in word2index.items():
|
| word_parts = mws_splitter.split(word)
|
|
|
| if len(word_parts) > 1:
|
|
|
| if not all(len(wp.strip()) > 0 for wp in word_parts):
|
| print(f"Found multiword with a blank component ({word_parts}). Skipping.")
|
| continue
|
|
|
| word_parts_indices = [word2index[wp] for wp in word_parts]
|
| wpi_docs = [set(word_corpus_index[wpi].keys()) for wpi in word_parts_indices]
|
| final_docs = wpi_docs[0]
|
| for wpid in wpi_docs[1:]:
|
| final_docs = final_docs.intersection(wpid)
|
|
|
|
|
| if len(final_docs) == 0:
|
| corpus_occurrences[index] = set()
|
| continue
|
|
|
| for doc_id in final_docs:
|
| wpis_positions = [word_corpus_index[wpi][doc_id] for wpi in word_parts_indices]
|
| for wpis_position in zip(*wpis_positions):
|
| diffs = [wpis_position[i + 1] - wpis_position[i] for i in range(len(wpis_position) - 1)]
|
| if set(diffs) == {1}:
|
| corpus_occurrences[index].add(doc_id)
|
| break
|
| else:
|
| corpus_occurrences[index] = set(word_corpus_index[index].keys())
|
|
|
| return corpus_occurrences
|
|
|
|
|
| def compute_cooccurrences(words_occurrences: Dict[int, Set[int]]) -> List[Tuple[str, str, int]]:
|
| cooccurrences = collections.Counter()
|
|
|
| doc2words = collections.defaultdict(list)
|
| for wid, doc_ids in words_occurrences.items():
|
| for doc_id in doc_ids:
|
| doc2words[doc_id].append(wid)
|
|
|
| for doc, words in tqdm(doc2words.items(), total=len(doc2words)):
|
| cooccurrences.update(itertools.combinations(words, 2))
|
|
|
| return [(w1, w2, coocs) for (w1, w2), coocs in cooccurrences.items()]
|
|
|
|
|
| def compute_pmi_files(
|
| vocabulary_path: str,
|
| corpus_path: str,
|
| output_dir: str,
|
| min_coocs_support: int,
|
| ) -> None:
|
| word2index, index2word = index_vocabulary(vocabulary_path)
|
| word_corpus_index = index_corpus(word2index, corpus_path)
|
| word_corpus_occurrences = compute_corpus_occurrences(word2index, word_corpus_index)
|
|
|
| words_occurrences = [(word, len(doc_ids)) for word, doc_ids in word_corpus_occurrences.items()]
|
| words_cooccurrences = compute_cooccurrences(word_corpus_occurrences)
|
|
|
| with open(f"{output_dir}/words_counter.tsv", "w") as f:
|
| for word, nocc in sorted(words_occurrences, key=lambda x: x[1], reverse=True):
|
| f.write(f"{index2word[word]}\t{nocc}\n")
|
|
|
| with open(f"{output_dir}/word_pairs_counter.tsv", "w") as f:
|
| for word1, word2, nocc in sorted(words_cooccurrences, key=lambda x: x[2], reverse=True):
|
| if nocc >= min_coocs_support:
|
| f.write(f"{index2word[word1]}\t{index2word[word2]}\t{nocc}\n")
|
|
|
|
|
| def parse_args() -> argparse.Namespace:
|
| parser = argparse.ArgumentParser()
|
| parser.add_argument("--vocabulary")
|
| parser.add_argument("--corpus")
|
| parser.add_argument("--output-dir")
|
| parser.add_argument("--min-coocs-support", type=int, default=1)
|
| return parser.parse_args()
|
|
|
|
|
| def main():
|
| args = parse_args()
|
| compute_pmi_files(args.vocabulary, args.corpus, args.output_dir, args.min_coocs_support)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
