init: app

.gitignore

@@ -0,0 +1 @@
+.python-version

app.py

@@ -0,0 +1,144 @@
+import gradio as gr
+import torch
+from transformers import GPT2LMHeadModel, T5Tokenizer
+
+model_name = "akiFQC/japanese-dialogpt-small-aozora"
+tokenizer = T5Tokenizer.from_pretrained(model_name)
+tokenizer.do_lower_case = True  # due to some bug of tokenizer config loading
+model = GPT2LMHeadModel.from_pretrained(model_name)
+
+
+class DialogGPT:
+    def __init__(self, tokenizer, model, n_candidate=4, param_lambda=0.1):
+        self.tokenizer = tokenizer
+        self.model = model
+        self.model.eval()
+        self.n_candidate = n_candidate
+        self.param_lambda = param_lambda
+
+    def _calc_single_scores(self, token_ids):
+        with torch.inference_mode():
+            candidate_token_ids = token_ids[:, :-1]
+            label_token_ids = token_ids[:, 1:]
+            outputs = self.model(candidate_token_ids, labels=label_token_ids)
+        _, logits = outputs[:2]
+        logits = torch.log_softmax(logits, dim=-1)
+
+        logit_at_target = logits.gather(
+            dim=-1, index=candidate_token_ids.unsqueeze(-1)
+        ).squeeze(-1)
+
+        # mask out pad token positio
+        mask_at_pad = candidate_token_ids == self.tokenizer.pad_token_id
+        # log_likelihood (b, l)
+        log_likelihood = logit_at_target
+        log_likelihood.masked_fill_(mask_at_pad, 0.0)
+        log_likelihood_per_candidate = log_likelihood.sum(dim=1)
+        # normalize by length
+        # log_likelihood_per_candidate = log_likelihood_per_candidate / (candidate_token_ids.shape[1] - mask_at_pad.sum(dim=1))
+        return log_likelihood_per_candidate
+
+    def _calc_scores(self, sequences, scores, input_ids=None):
+        transition_scores = model.compute_transition_scores(
+            sequences, scores, normalize_logits=True
+        )
+        if input_ids is None:
+            input_length = 0
+        else:
+            input_length = input_ids.shape[1]
+        generated_tokens = sequences[:, input_length:]  # n x l
+        assert (
+            generated_tokens.shape[1] == transition_scores.shape[1]
+        ), f"{generated_tokens.shape[1]} != {transition_scores.shape[1]}"
+        # print(transition_scores.shape)
+        # print(generated_tokens)
+        transition_scores.masked_fill_(
+            generated_tokens == self.tokenizer.pad_token_id, 0.0
+        )
+        transition_scores = transition_scores.sum(dim=1)
+        # print(transition_scores)
+        return transition_scores
+
+    def reply(self, reply, history) -> str:
+        chat_history_ids = torch.LongTensor(history).unsqueeze(0)
+        # encode the new user input, add the eos_token and return a tensor in Pytorch
+        new_user_input_ids = self.tokenizer.encode(
+            reply + self.tokenizer.eos_token, return_tensors="pt"
+        )
+
+        # append the new user input tokens to the chat history
+        bot_input_ids = (
+            torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
+            if chat_history_ids is not None
+            else new_user_input_ids
+        )
+
+        # generated a response while limiting the total chat history to 1000 tokens,
+        with torch.inference_mode():
+            output = model.generate(
+                bot_input_ids,
+                pad_token_id=self.tokenizer.pad_token_id,
+                do_sample=True,
+                top_p=0.93,
+                temperature=0.5,
+                repetition_penalty=1.17,
+                max_time=10,
+                num_return_sequences=self.n_candidate,
+                max_length=512,
+                min_length=2,
+                forced_eos_token_id=self.tokenizer.pad_token_id,
+                return_dict_in_generate=True,
+                output_scores=True,
+                min_new_tokens=2,
+            )
+
+        # score of each candidate
+        scores_condition_s2t = self._calc_scores(
+            sequences=output.sequences, scores=output.scores, input_ids=bot_input_ids
+        )
+        new_token_ids = output.sequences[:, bot_input_ids.shape[-1] :]
+        single_scores = self._calc_single_scores(new_token_ids) * self.param_lambda
+
+        total_scores = scores_condition_s2t - single_scores
+        id_selected = torch.argmax(total_scores)
+
+        chat_history_ids = output.sequences[id_selected].unsqueeze(
+            0
+        )  # update chat history
+        # remove pad token
+        chat_history_ids = chat_history_ids[
+            :, chat_history_ids[0] != self.tokenizer.pad_token_id
+        ]
+        replay_string = tokenizer.decode(
+            chat_history_ids[:, :][0], skip_special_tokens=False
+        )
+        return replay_string, chat_history_ids[0].tolist()
+
+
+bot = DialogGPT(
+    tokenizer,
+    model,
+)
+
+
+def predict(input, history=[]):
+    replay_string, history = bot.reply(input, history)
+    response = replay_string.split(tokenizer.eos_token)
+    response = [
+        (response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)
+    ]  # convert to tuples of list
+    return response, history
+
+
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot()
+    state = gr.State([])
+
+    with gr.Row():
+        txt = gr.Textbox(
+            show_label=False, placeholder="Enter text and press enter"
+        ).style(container=False)
+
+    txt.submit(predict, [txt, state], [chatbot, state])
+
+demo.launch()