<reponame>PartiallyTyped/tydiqa<filename>snlp/prepare/model.py
import enum
from functools import partial, reduce
from itertools import chain, combinations, count, starmap
from pathlib import Path
from gensim.models import rpmodel
from numpy.core.fromnumeric import shape
from numpy.lib.arraysetops import isin

from scipy import sparse
import snlp

import gensim
from snlp.structs import CompressedCorpus, CompressedMinSpanSample, MinSpanEntry, ProcessedSample, TokenizedMinSpanSample
from typing import Dict, Iterable, List, Optional, Tuple
from tqdm.std import tqdm, trange
from snlp.util import argparsers as parsers, dumpers
from snlp.util import loaders
from gensim.utils import pickle, tokenize as gtokensize
from nltk import sent_tokenize
from gensim.parsing import preprocessing
from gensim import models
from snlp.prepare.vectorise import collect_text
import pickle as pkl
import numpy as np
from  sklearn.metrics.pairwise import cosine_similarity
def corpus_to_sparse_matrix(corpus, columns):
    """
    takes in corpus and returns sparse matrix
    """
    # create empty sparse matrix
    sparse_matrix = sparse.lil_matrix((len(corpus), columns))
    # fill in sparse matrix
    for i in trange(len(corpus), desc="unsparcifying"):
        for j in range(len(corpus[i])):
            if corpus[i][j][0] >= columns:
                continue
            sparse_matrix[i, corpus[i][j][0]] = corpus[i][j][1]
    return sparse_matrix

def statistical(data: List[CompressedMinSpanSample], model:gensim.interfaces.TransformationABC, output: str, vector_size:int=500):
    _, questions, passages, answers, ids = zip(*data)

    passage_features = list(tqdm(model[passages], desc="Computing passage features"))
    question_features = list(tqdm(model[questions], desc="Computing question features"))
    passage_features = corpus_to_sparse_matrix(passage_features, vector_size)
    question_features = corpus_to_sparse_matrix(question_features, vector_size)
    features = sparse.hstack([passage_features, question_features]).toarray()
    
    data = list(starmap(ProcessedSample, zip(features, answers, ids)))
    dumpers.processed(data, output)

def embeddings(data: List[CompressedMinSpanSample], model:gensim.models.Word2Vec, output: str, vector_size=128):
    _, questions, passages, answers, ids = zip(*data)
    def sentence_embedding(sentence):
        embd = []
        for word in sentence:
            assert isinstance(word, str), word
            if word not in model.wv:
                continue
            embd.append(model.wv[word])
        embd = np.asarray(embd)
        if embd.shape[0] == 0:
            embd = np.zeros((1, model.vector_size))
        embd = embd.mean(axis=0)
        return embd
    model.layer1_size

    features = []
    for question, passage in zip(questions, tqdm(passages, "Computing Features", leave=False)):
        q_embd = sentence_embedding(question)
        p_embd = []
        for n_sentences in range(1, min(3, len(passage))):
            for sents in combinations(passage, n_sentences):
                w = []
                for sent in sents:
                    w.extend(sent)
                p_embd.append(sentence_embedding(sent))
        p_embd = np.asarray(p_embd)
        cos_sim = cosine_similarity(p_embd, q_embd.reshape(1, -1)).ravel()
        cos_sim = np.sort(cos_sim)
        passage_features = np.zeros(vector_size)
        np.put(passage_features, np.arange(vector_size), cos_sim, mode="clip")
        features.append(passage_features)
    features = np.asanyarray(features)
    data = list(starmap(ProcessedSample, zip(features, answers, ids)))
    dumpers.processed(data, output)

def load_model(path):
    model = gensim.utils.SaveLoad.load(path)
    model = type(model).load(path)
    return model

def argument_parser():
    mods = [("embeddings",embeddings), ("statistical",statistical)]
    mods = {
        name:parsers.generate_parser(
            mod,
            dest="main",
            add_run=True,
            data=loaders.load_compressed,
            model=load_model,
        )
        for name, mod in mods
    }
    
    parser = parsers.no_help()
    mod_parsers = parser.add_subparsers()
    for mod, m_parser in mods.items():
        mod_parsers.add_parser(
            mod,
            parents=[m_parser]
        )

    return parser

import enum
from functools import partial, reduce
from itertools import chain, combinations, count, starmap
from pathlib import Path
from gensim.models import rpmodel
from numpy.core.fromnumeric import shape
from numpy.lib.arraysetops import isin

from scipy import sparse
import snlp

import gensim
from snlp.structs import CompressedCorpus, CompressedMinSpanSample, MinSpanEntry, ProcessedSample, TokenizedMinSpanSample
from typing import Dict, Iterable, List, Optional, Tuple
from tqdm.std import tqdm, trange
from snlp.util import argparsers as parsers, dumpers
from snlp.util import loaders
from gensim.utils import pickle, tokenize as gtokensize
from nltk import sent_tokenize
from gensim.parsing import preprocessing
from gensim import models
from snlp.prepare.vectorise import collect_text
import pickle as pkl
import numpy as np
from  sklearn.metrics.pairwise import cosine_similarity
def corpus_to_sparse_matrix(corpus, columns):
    """
    takes in corpus and returns sparse matrix
    """
    # create empty sparse matrix
    sparse_matrix = sparse.lil_matrix((len(corpus), columns))
    # fill in sparse matrix
    for i in trange(len(corpus), desc="unsparcifying"):
        for j in range(len(corpus[i])):
            if corpus[i][j][0] >= columns:
                continue
            sparse_matrix[i, corpus[i][j][0]] = corpus[i][j][1]
    return sparse_matrix

def statistical(data: List[CompressedMinSpanSample], model:gensim.interfaces.TransformationABC, output: str, vector_size:int=500):
    _, questions, passages, answers, ids = zip(*data)

    passage_features = list(tqdm(model[passages], desc="Computing passage features"))
    question_features = list(tqdm(model[questions], desc="Computing question features"))
    passage_features = corpus_to_sparse_matrix(passage_features, vector_size)
    question_features = corpus_to_sparse_matrix(question_features, vector_size)
    features = sparse.hstack([passage_features, question_features]).toarray()
    
    data = list(starmap(ProcessedSample, zip(features, answers, ids)))
    dumpers.processed(data, output)

def embeddings(data: List[CompressedMinSpanSample], model:gensim.models.Word2Vec, output: str, vector_size=128):
    _, questions, passages, answers, ids = zip(*data)
    def sentence_embedding(sentence):
        embd = []
        for word in sentence:
            assert isinstance(word, str), word
            if word not in model.wv:
                continue
            embd.append(model.wv[word])
        embd = np.asarray(embd)
        if embd.shape[0] == 0:
            embd = np.zeros((1, model.vector_size))
        embd = embd.mean(axis=0)
        return embd
    model.layer1_size

    features = []
    for question, passage in zip(questions, tqdm(passages, "Computing Features", leave=False)):
        q_embd = sentence_embedding(question)
        p_embd = []
        for n_sentences in range(1, min(3, len(passage))):
            for sents in combinations(passage, n_sentences):
                w = []
                for sent in sents:
                    w.extend(sent)
                p_embd.append(sentence_embedding(sent))
        p_embd = np.asarray(p_embd)
        cos_sim = cosine_similarity(p_embd, q_embd.reshape(1, -1)).ravel()
        cos_sim = np.sort(cos_sim)
        passage_features = np.zeros(vector_size)
        np.put(passage_features, np.arange(vector_size), cos_sim, mode="clip")
        features.append(passage_features)
    features = np.asanyarray(features)
    data = list(starmap(ProcessedSample, zip(features, answers, ids)))
    dumpers.processed(data, output)

def load_model(path):
    model = gensim.utils.SaveLoad.load(path)
    model = type(model).load(path)
    return model

def argument_parser():
    mods = [("embeddings",embeddings), ("statistical",statistical)]
    mods = {
        name:parsers.generate_parser(
            mod,
            dest="main",
            add_run=True,
            data=loaders.load_compressed,
            model=load_model,
        )
        for name, mod in mods
    }
    
    parser = parsers.no_help()
    mod_parsers = parser.add_subparsers()
    for mod, m_parser in mods.items():
        mod_parsers.add_parser(
            mod,
            parents=[m_parser]
        )

    return parser
takes in corpus and returns sparse matrix # create empty sparse matrix # fill in sparse matrix