fasttext_fsc.py

from __future__ import annotations
import torch
from torchtyping import TensorType
from .fasttext_jp_embedding import FastTextJpModel, FastTextJpConfig
from transformers.modeling_outputs import SequenceClassifierOutput


class FastTextForSeuqenceClassificationConfig(FastTextJpConfig):
 """FastTextJpModelのConfig
 """
 model_type = "fasttext_classification"

 def __init__(self,
 ngram: int | list[int] = 2,
 tokenizer_class="FastTextJpTokenizer",
 **kwargs):
 """初期化処理

 Args:
 ngram (int | list[int], optional):
 文章を分割する際のNgram。
 tokenizer_class (str, optional): 
 tokenizer_classを指定しないと、pipelineから読み込まれません。
 config.jsonに記載されます。
 """
 if isinstance(ngram, int):
 self.ngrams = [ngram]
 elif isinstance(ngram, list):
 self.ngrams = ngram
 else:
 raise TypeError(f"got unknown type {type(ngram)}")
 kwargs["tokenizer_class"] = tokenizer_class
 super().__init__(**kwargs)


class NgramForSeuqenceClassification():

 def __init__(self):
 ...

 def __call__(self, sentence: TensorType["A", "vectors"],
 candidate_label: TensorType["B", "vectors"],
 ngram: int) -> TensorType[3]:
 """Ngramで文章を分けてコサイン類似度を算出する。

 Args:
 sentence (TensorType["A", "vectors"]): 文章ベクトル
 candidate_label (TensorType["B", "vectors"]): ラベルベクトル
 ngram (int): Ngram

 Returns:
 TensorType[3]:
 文章の類似度。[Entailment, Neutral, Contradiction]
 """

 sentence_ngrams = self.split_ngram(sentence, ngram)

 candidate_label_mean = torch.mean(candidate_label, dim=0, keepdim=True)
 p = self.cosine_similarity(sentence_ngrams, candidate_label_mean)
 return torch.tensor([torch.log(p), -torch.inf, torch.log(1 - p)])

 def cosine_similarity(
 self, sentence_ngrams: TensorType["ngrams", "vectors"],
 candidate_label_mean: TensorType[1, "vectors"]) -> TensorType[1]:
 """コサイン類似度を計算する。

 Args:
 sentence_ngrams (TensorType["ngrams", "vectors"]):
 Ngram化された文章ベクトル
 candidate_label_mean (TensorType[1, "vectors"]): 
 ラベルベクトル

 Returns:
 TensorType[1]: _description_
 """

 res = torch.tensor(0.)
 for i in range(len(sentence_ngrams)):
 sw = sentence_ngrams[i]
 p = torch.nn.functional.cosine_similarity(sw,
 candidate_label_mean[0],
 dim=0)
 if p > res:
 res = p
 return res

 def split_ngram(self, sentences: TensorType["A", "vectors"],
 n: int) -> TensorType["ngrams", "vectors"]:
 """AとBの関連度を計算します。
 Args:
 sentences(TensorType["A", "vectors"]): 
 対象の文章
 n(int):
 ngram
 Returns:
 TensorType["ngrams", "vectors"]:
 Ngram化された文章
 """

 res = []
 if len(sentences) <= n:
 return torch.stack([torch.mean(sentences, dim=0, keepdim=False)])
 for i in range(len(sentences) - n + 1):
 ngram = sentences[i:i + n]
 res.append(torch.mean(ngram, dim=0, keepdim=False))
 return torch.stack(res)


class NgramsForSeuqenceClassification():

 def __init__(self, config: FastTextForSeuqenceClassificationConfig):
 self.max_ngrams = config.ngrams
 self.ngram_layer = NgramForSeuqenceClassification()

 def __call__(self, sentence: TensorType["A", "vectors"],
 candidate_label: TensorType["B", "vectors"]) -> TensorType[3]:
 """AとBの関連度を計算します。
 Args:
 sentence(TensorType["A", "vectors"]): 
 対象の文章
 candidate_label(TensorType["B", "vectors"]):
 ラベルの文章

 Returns:
 TensorType[3]:
 文章の類似度。[Entailment, Neutral, Contradiction]
 """

 res = [-torch.inf, -torch.inf, -torch.inf]
 for ngram in self.max_ngrams:
 logit = self.ngram_layer(sentence, candidate_label, ngram)
 if logit[0] > res[0]:
 res = logit
 return torch.tensor(res)


class BatchedNgramsForSeuqenceClassification():

 def __init__(self, config: FastTextForSeuqenceClassificationConfig):
 self.ngrams_layer = NgramsForSeuqenceClassification(config)

 def __call__(
 self,
 last_hidden_state: TensorType["batch", "A+B", "vectors"],
 token_type_ids: TensorType["batch", "A+B"],
 attention_mask: TensorType["batch", "A+B"],
 ) -> TensorType["batch", 3]:
 """AとBの関連度を計算します。
 Args:
 last_hidden_state(TensorType["batch", "A+B", "vectors"]): 
 embeddingsの値。
 token_type_ids(TensorType["A+B"]):
 文章のid。0か1で、Bの場合1。
 attention_mask(TensorType["A+B"]):
 padを識別する。0か1で、padの場合1。

 Returns:
 TensorType["batch", 3]:
 文章の類似度。[Entailment, Neutral, Contradiction]
 """

 logits = []
 embeddings = last_hidden_state
 for idx in range(len(embeddings)):
 vec = embeddings[idx]
 # token_type_ids == 0が文章、1がラベルです。
 token_type_ids = token_type_ids[idx]
 # attention_mask == 1がパディングでないもの
 attention_mask = attention_mask[idx]

 sentence, candidate_label = self.split_sentence(
 vec, token_type_ids, attention_mask)
 logit = self.ngrams_layer(sentence, candidate_label)
 logits.append(logit)
 logits = torch.tensor(logits)
 return logits

 def split_sentence(
 self, vec: TensorType["A+B", "vectors"],
 token_type_ids: TensorType["A+B"], attention_mask: TensorType["A+B"]
 ) -> tuple[TensorType["A", "vectors"], TensorType["B", "vectors"]]:
 """CrossEncoderになっているので、文章を分割します。
 
 Args:
 vec(TensorType["A+B","vectors"]):
 単語ベクトル

 token_type_ids(TensorType["A+B"]):
 文章のid。0か1で、Bの場合1。

 attention_mask(TensorType["A+B"]):
 padを識別する。0か1で、padの場合1。

 Returns:
 tuple[TensorType["A", "vectors"], TensorType["B", "vectors"]]:
 AとBの文章を分割して返します。
 """

 sentence = vec[torch.logical_and(token_type_ids == 0,
 attention_mask == 1)]
 candidate_label = vec[torch.logical_and(token_type_ids == 1,
 attention_mask == 1)]
 return sentence, candidate_label


class FastTextForSeuqenceClassification(FastTextJpModel):
 """FastTextのベクトルをベースとした分類を行います。
 """

 def __init__(self, config: FastTextForSeuqenceClassificationConfig):

 self.layer = BatchedNgramsForSeuqenceClassification(config)
 super().__init__(config)

 def forward(
 self,
 input_ids: TensorType["batch", "A+B", "vecotors"] = None,
 attention_mask: TensorType["batch", "A+B"] = None,
 token_type_ids: TensorType["batch", "A+B"] = None
 ) -> SequenceClassifierOutput:
 """候補となるラベルから分類を行います。

 Returns:
 SequenceClassifierOutput: 候補が正解している確率
 """
 outputs = self.word_embeddings(input_ids)
 logits = self.layer(last_hidden_state=outputs,
 attention_mask=attention_mask,
 token_type_ids=token_type_ids)

 return SequenceClassifierOutput(
 loss=None,
 logits=logits,
 hidden_states=None,
 attentions=None,
 )


# AutoModelに登録が必要だが、いろいろやり方が変わっているようで定まっていない。(2022/11/6)
# https://huggingface.co/docs/transformers/custom_models#sending-the-code-to-the-hub
FastTextForSeuqenceClassificationConfig.register_for_auto_class()
FastTextForSeuqenceClassification.register_for_auto_class(
 "AutoModelForSequenceClassification")