# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # variant2vec
#
# Learn an emedding of variants annotations (HGVS, BIC etc...), Hugo gene symbols, and classifications (pathogenic, benign...) from the text of pubmed articles.
#
# Goal to see if something like [this](https://www.nature.com/articles/s41586-019-1335-8) paper can be applied to determine characteristics of a variant, or at lease identify the gene associated with the variant from the paper where multiple variants may be mentioned with no canonical gene or transcript qualifier.

# +
import os
import sys
import re
import sqlite3
import pandas as pd
import gzip
import glob
import datetime

def log(text):
    # Print to jupyter notebook and console so we can monitor progress in k8s
    print(text)
    sys.__stdout__.write(text + "\n")
    sys.__stdout__.flush()
    
# Syntatic sugar for debug vs. train parameters
def debug(debug_param, no_debug_param):
    return debug_param if os.environ.get("DEBUG") == "True" else no_debug_param
log(debug("DEBUG: ON", "DEBUG: OFF"))

crawl_path = "/public/groups/brcaexchange/literature-search/crawl"
# -

# ## Ingest
# Load the text from the pubMunch [pubStore](https://github.com/maximilianh/pubMunch/blob/master/lib/pubStore.py) generated in a BRCA Exchange literature search crawl

connection = sqlite3.connect(f"file:{crawl_path}/text/articles.db?mode=ro", uri=True)
articles = pd.read_sql_query("SELECT * FROM articles", connection)
articles.pmid = articles.pmid.astype(str)
print("{} articles loaded from the articles sqlite database".format(articles.shape[0]))
articles.head()

# +
# %%time
log("Loading documents...")
articles = pd.DataFrame()
for path in glob.glob(f"{crawl_path}/text/*.files.gz"):
    articles = articles.append(pd.read_csv(path, sep="\t", encoding="utf-8"))
    
print("{} article text blocks consuming {:.2f} megabytes".format(
    articles.shape[0], articles.memory_usage(index=True).sum() / 2**20))
# -

articles.head()

# %%time
variants = set()
pattern = re.compile("c\.\d+[atcgATCG]+>[atcgATCG]+")
for article in articles.itertuples():
    variants = variants.union(set(re.findall(pattern, article.content)))
print("Found {} variants".format(len(variants)))

# ## Clean
# https://machinelearningmastery.com/clean-text-machine-learning-python/

# +
# !pip3 install --user --upgrade --quiet unidecode
import unidecode

# !pip3 install --user --upgrade --quiet nltk
import nltk
import nltk.tokenize
import nltk.stem.porter

nltk.download("stopwords")
stop_words = nltk.corpus.stopwords.words('english')

stemmer = nltk.stem.porter.PorterStemmer()

# +
log("Cleaning...")
word_regex = re.compile("\s+|[,\(\)]")

def clean(text):
    # Convert to ascii and replace returns with spaces and lower case
    text = re.sub(r"[\x07]+", ' ', unidecode.unidecode(text)).lower()
    
    # split into words
#     words = nltk.tokenize.word_tokenize(text)
    
    words = word_regex.split(text)
    
    # Drop all single character words and stop words - effectively removes punctuation
    words = [w for w in words if len(w) > 3 and w not in stop_words]
    
    # Stem (lowercase's as a side affect as well)
    words = [stemmer.stem(w) for w in words]
    
    return words

print(clean(articles.iloc[42].content[0:2000]))
# -

# %%time
texts = [clean(article.content) for article in articles.iloc[0:None:debug(100,None)].itertuples()]
print("Cleaned {} texts".format(len(texts)))

# +
print("Before:")
print(articles.iloc[2].content[0:2000])

print("After:")
print(texts[2][0:250])
# -

# ## Train
#
# References:
#
# https://radimrehurek.com/gensim/tut1.html
#
# https://machinelearningmastery.com/develop-word-embeddings-python-gensim/

# +
# !pip3 install --user --upgrade --quiet gensim

import gensim
import gensim.models
# -

# model = gensim.models.Word2Vec(documents, size=150, window=10, min_count=2, workers=10, iter=10)
log("Training...")
model = gensim.models.Word2Vec(texts, max_vocab_size=5000)

# +
# Timestamp for output files
log("Saving...")
timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
print("Training run at {}".format(timestamp))

path = os.path.expanduser("~/data/variante2vec/{}.model").format(debug("debug", timestamp))
print("Saved to:",path)

model.save(path)
model = gensim.models.Word2Vec.load(path)

StarcoderdataJupyterScriptsDedupFiltered