<commit_before>#!/usr/bin/env python3

import argparse
import concurrent.futures
import csv
import itertools
import sys
from collections import defaultdict, Counter
from signal import signal, SIGINT

from sklearn.feature_extraction import DictVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.metrics.pairwise import cosine_similarity as sim
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import Binarizer

signal(SIGINT, lambda signum, frame: sys.exit(1))

WEIGHT = {
    'tf': TfidfTransformer(use_idf=False),
    'idf': Pipeline([('binary', Binarizer()), ('idf', TfidfTransformer())]),
    'tfidf': TfidfTransformer()
}

parser = argparse.ArgumentParser()
parser.add_argument('--synsets', required=True, type=argparse.FileType('r'))
parser.add_argument('--isas', required=True, type=argparse.FileType('r'))
parser.add_argument('--weight', choices=WEIGHT.keys(), default='tfidf')
parser.add_argument('-k', nargs='?', type=int, default=6)
args = parser.parse_args()

synsets, index, lexicon = {}, defaultdict(list), set()

with args.synsets as f:
    reader = csv.reader(f, delimiter='\t', quoting=csv.QUOTE_NONE)

    for row in reader:
        synsets[row[0]] = [word for word in row[2].split(', ') if word]

        for word in synsets[row[0]]:
            index[word].append(row[0])

        lexicon.update(synsets[row[0]])

index = {word: {id: i + 1 for i, id in enumerate(ids)} for word, ids in index.items()}

isas = defaultdict(set)

with args.isas as f:
    reader = csv.reader(f, delimiter='\t', quoting=csv.QUOTE_NONE)

    for row in reader:
        if len(row) > 1 and row[0] in lexicon and row[1] in lexicon:
            isas[row[0]].add(row[1])

hctx = {}

for id, words in synsets.items():
    hypernyms = Counter(itertools.chain(*(isas[word] for word in words if word in isas)))

    if hypernyms:
        hctx[id] = hypernyms

v = Pipeline([('dict', DictVectorizer()), (args.weight, WEIGHT[args.weight])]).fit(hctx.values())


def emit(id):
    if not id in hctx:
        return (id, {})

    hvector, candidates = v.transform(hctx[id]), Counter()

    for hypernym in hctx[id]:
        hsenses = Counter({hid: sim(v.transform(Counter(synsets[hid])), hvector).item(0) for hid in index[hypernym]})

        for hid, cosine in hsenses.most_common(1):
            if cosine > 0:
                candidates[(hypernym, hid)] = cosine<commit_msg>Stop using the csv module in link.py<commit_after>#!/usr/bin/env python3

import argparse
import concurrent.futures
import itertools
import sys
from collections import defaultdict, Counter
from signal import signal, SIGINT

from sklearn.feature_extraction import DictVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.metrics.pairwise import cosine_similarity as sim
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import Binarizer

signal(SIGINT, lambda signum, frame: sys.exit(1))

WEIGHT = {
    'tf': TfidfTransformer(use_idf=False),
    'idf': Pipeline([('binary', Binarizer()), ('idf', TfidfTransformer())]),
    'tfidf': TfidfTransformer()
}

parser = argparse.ArgumentParser()
parser.add_argument('--synsets', required=True, type=argparse.FileType('r'))
parser.add_argument('--isas', required=True, type=argparse.FileType('r'))
parser.add_argument('--weight', choices=WEIGHT.keys(), default='tfidf')
parser.add_argument('-k', nargs='?', type=int, default=6)
args = parser.parse_args()

synsets, index, lexicon = {}, defaultdict(list), set()

for line in args.synsets:
    line = line.rstrip()

    if not line:
        continue

    row = line.split('\t')

    synsets[row[0]] = [word for word in row[2].split(', ') if word]

    for word in synsets[row[0]]:
        index[word].append(row[0])

    lexicon.update(synsets[row[0]])

index = {word: {id: i + 1 for i, id in enumerate(ids)} for word, ids in index.items()}

isas = defaultdict(set)

for line in args.isas:
    line = line.rstrip()

    if not line:
        continue

    row = line.split('\t')

    if len(row) > 1 and row[0] in lexicon and row[1] in lexicon:
        isas[row[0]].add(row[1])

hctx = {}

for id, words in synsets.items():
    hypernyms = Counter(itertools.chain(*(isas[word] for word in words if word in isas)))

    if hypernyms:
        hctx[id] = hypernyms

v = Pipeline([('dict', DictVectorizer()), (args.weight, WEIGHT[args.weight])]).fit(hctx.values())


def emit(id):
    if not id in hctx:
        return (id, {})

    hvector, candidates = v.transform(hctx[id]), Counter()

    for hypernym in hctx[id]:
        hsenses = Counter({hid: sim(v.transform(Counter(synsets[hid])), hvector).item(0) for hid in index[hypernym]})

        for hid, cosine in hsenses.most_common(1):
            if cosine > 0:
                candidates[(hypernym, hid)] = cosine
{'fraction_non_alphanumeric': 0.14526184538653367, 'fraction_numerical': 0.007065669160432253, 'mean_word_length': 4.045073375262055, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6896068', 'n_tokens_mistral': 1676, 'n_tokens_neox': 1631, 'n_words': 421}