File size: 4,760 Bytes
6dc7c27 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | 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)
# the word parts does not appear altogether
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()
|