util.py

import os
import re
import shutil
import tempfile
from pathlib import Path
from typing import List

import flair
import gensim
import numpy as np
import requests
import torch
from flair.data import Sentence
from flair.embeddings import StackedEmbeddings, PooledFlairEmbeddings, BytePairEmbeddings, \
    TokenEmbeddings
from flair.file_utils import Tqdm


def print_flag(content, big=True):
    l = len(content)
    bar = '#' * (8 + l)
    if big:
        print(f'{bar}\n'
              f'### {content} ###\n'
              f'{bar}')
    else:
        print(f'### {content} ###')


def get_from_cache(url: str, filename: str, cache_dir: Path = None) -> Path:
    """
    Given a URL, look for the corresponding dataset in the local cache.
    If it's not there, download it. Then return the path to the cached file.
    """
    cache_dir.mkdir(parents=True, exist_ok=True)

    # get cache path to put the file
    cache_path = cache_dir / filename
    if cache_path.exists():
        return cache_path

    # make HEAD request to check ETag
    response = requests.head(url, headers={"User-Agent": "Flair"})
    if response.status_code != 200:
        raise IOError(
            f"HEAD request failed for url {url} with status code {response.status_code}."
        )

    # add ETag to filename if it exists
    # etag = response.headers.get("ETag")

    if not cache_path.exists():
        # Download to temporary file, then copy to cache dir once finished.
        # Otherwise you get corrupt cache entries if the download gets interrupted.
        fd, temp_filename = tempfile.mkstemp()
        flair.logger.info("%s not found in cache, downloading to %s", url, temp_filename)

        # GET file object
        req = requests.get(url, stream=True, headers={"User-Agent": "Flair"})
        content_length = req.headers.get("Content-Length")
        total = int(content_length) if content_length is not None else None
        progress = Tqdm.tqdm(unit="B", total=total)
        with open(temp_filename, "wb") as temp_file:
            for chunk in req.iter_content(chunk_size=1024):
                if chunk:  # filter out keep-alive new chunks
                    progress.update(len(chunk))
                    temp_file.write(chunk)

        progress.close()

        flair.logger.info("copying %s to cache at %s", temp_filename, cache_path)
        shutil.copyfile(temp_filename, str(cache_path))
        flair.logger.info("removing temp file %s", temp_filename)
        os.close(fd)
        os.remove(temp_filename)

    return cache_path


class SpanishHealthCorpusEmbeddings(TokenEmbeddings):
    def __init__(self, embeddings: str, field: str = None, cache_dir=Path(flair.cache_root) / "embeddings"):
        self.name = embeddings

        if embeddings not in ('wang2vec', 'fastText'):
            raise ValueError(f"Embedding name must be in {('wang2vec', 'fastText')}!")

        api_url = "http://service.hucompute.org/embeddings/api/v1/embeddings"
        vector_file = f"es_health_{embeddings}.vec"
        vector_url = f"{api_url}/SpanishHealthCorpus_form_{vector_file}/download"
        embeddings = cache_dir.joinpath(vector_file)
        if not cache_dir.joinpath(vector_file).exists():
            get_from_cache(vector_url, vector_file, cache_dir=cache_dir)

        if str(embeddings).endswith(".vec"):
            self.precomputed_word_embeddings = gensim.models.KeyedVectors.load_word2vec_format(
                str(embeddings), binary=False, unicode_errors='replace'
            )

        self.name: str = str(embeddings)
        self.static_embeddings = True

        self.field = field

        self.__embedding_length: int = self.precomputed_word_embeddings.vector_size
        super().__init__()

    @property
    def embedding_length(self) -> int:
        return self.__embedding_length

    def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:

        for i, sentence in enumerate(sentences):

            for token, token_idx in zip(sentence.tokens, range(len(sentence.tokens))):

                if "field" not in self.__dict__ or self.field is None:
                    word = token.text
                else:
                    word = token.get_tag(self.field).value

                if word in self.precomputed_word_embeddings:
                    word_embedding = self.precomputed_word_embeddings[word]
                elif word.lower() in self.precomputed_word_embeddings:
                    word_embedding = self.precomputed_word_embeddings[word.lower()]
                elif (
                        re.sub(r"\d", "#", word.lower()) in self.precomputed_word_embeddings
                ):
                    word_embedding = self.precomputed_word_embeddings[
                        re.sub(r"\d", "#", word.lower())
                    ]
                elif (
                        re.sub(r"\d", "0", word.lower()) in self.precomputed_word_embeddings
                ):
                    word_embedding = self.precomputed_word_embeddings[
                        re.sub(r"\d", "0", word.lower())
                    ]
                else:
                    word_embedding = np.zeros(self.embedding_length, dtype="float")

                word_embedding = torch.FloatTensor(word_embedding)

                token.set_embedding(self.name, word_embedding)

        return sentences

    def __str__(self):
        return self.name

    def extra_repr(self):
        # fix serialized models
        if "embeddings" not in self.__dict__:
            self.embeddings = self.name

        return f"'{self.embeddings}'"


def get_embeddings(pooling_op='min'):
    return StackedEmbeddings(embeddings=[
        # pre-trained embeddings
        PooledFlairEmbeddings(
            'es-forward', pooling=pooling_op,
        ),
        PooledFlairEmbeddings(
            'es-backward', pooling=pooling_op,
        ),
        BytePairEmbeddings(
            language='es', dim=300,
        ),

        # self-trained embeddings
        SpanishHealthCorpusEmbeddings('wang2vec'),
        # SpanishHealthCorpusEmbeddings('fastText'),
    ])