Create T5Mixer

app.py

@@ -1,4 +1,5 @@
 from dummymixer import DummyMixer
+from t5mixer import T5Mixer
 from frontend import launch
 
-launch(DummyMixer())
+launch(T5Mixer())

t5mixer.py

@@ -0,0 +1,52 @@
+from mixer import Mixer
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+
+class T5Mixer(Mixer):
+    def __init__(self) -> None:
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            "llm-book/t5-base-long-livedoor-news-corpus")
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(
+            "llm-book/t5-base-long-livedoor-news-corpus")
+        self.tokenid_to_tokentext = {i: t for t,
+                                     i in self.tokenizer.get_vocab().items()}
+
+    def get_encoder_state(self, sentence: str):
+        inputs = self.tokenizer(sentence, return_tensors="pt")
+        eo = self.model.encoder.forward(**inputs)
+        es = eo["last_hidden_state"]
+        return es
+
+    def get_mixed_encode_state(
+            self, sentence_A: str, sentence_B: str, A_ratio: float = 0.5):
+        es_A = self.get_encoder_state(sentence_A)
+        es_B = self.get_encoder_state(sentence_B)
+        n_tokens_A = es_A.size(1)
+        n_tokens_B = es_B.size(1)
+        if n_tokens_A >= n_tokens_B:
+            es = es_A.clone().detach()*A_ratio
+            es[:, :n_tokens_B, :] += es_B*(1.-A_ratio)
+        else:
+            es = es_B.clone().detach()*(1.-A_ratio)
+            es[:, :n_tokens_A, :] += es_A*A_ratio
+        return es
+
+    def mix_sentences(self, sentence_A: str, sentence_B: str, A_ratio: float, max_n_tokens: int = 140):
+        es = self.get_mixed_encode_state(sentence_A, sentence_B, A_ratio)
+        to = torch.tensor([[self.tokenizer.pad_token_id]])
+        for i in range(max_n_tokens):
+            od = self.model.decoder.forward(
+                input_ids=to,
+                encoder_hidden_states=es+torch.randn_like(es)*noise_rate
+            )
+            sd = od.last_hidden_state
+            l = self.model.lm_head(sd[0, -1, :])
+            t_next = l.argmax()
+            ttext_next = self.tokenid_to_tokentext[int(t_next)]
+            to = torch.cat((to, t_next[None, None]), dim=-1)
+            if t_next == self.tokenizer.eos_token_id:
+                break
+        sentence = self.tokenizer.batch_decode(to)[0]
+        return sentence