src/chatbots/gptjbot.py

from transformers import GPTJForCausalLM, AutoTokenizer
import os
import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class DialoGPT:
def __init__(
self,
model_name: str = "EleutherAI/gpt-j-6B",
local_path="./models/gpt-j-6B",
):
if not os.path.exists(local_path):
GPTJForCausalLM.from_pretrained(model_name).save_pretrained(
local_path,
revision="float16",
torch_dtype=torch.float16,
)
AutoTokenizer.from_pretrained(model_name).save_pretrained(local_path)

self.model = GPTJForCausalLM.from_pretrained(
local_path,
revision="float16",
torch_dtype=torch.float16,
)
self.tokenizer = AutoTokenizer.from_pretrained(local_path)

def __call__(self, inputs: str) -> str:
input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids.to(device)
generated_ids = self.model.to(device).generate(
input_ids, do_sample=True, temperature=0.9, max_length=200
)
generated_text = self.tokenizer.decode(generated_ids[0])

return generated_text

def run(self):
while True:
user_input = input("User: ")
print("Bot:", self(user_input))


if __name__ == "__main__":
bot = DialoGPT()
bot.run()

src/chatbots/dialogpt.py

@@ -0,0 +1,33 @@
+"""
+Adapted from:
+https://www.machinecurve.com/index.php/2021/03/16/easy-chatbot-with-dialogpt-machine-learning-and-huggingface-transformers/
+"""
+
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import os
+
+
+class DialoGPT:
+    def __init__(
+        self,
+        model_name: str ='microsoft/DialoGPT-large',
+    ):
+        if not os.path.exists('./models/dialogpt'):
+            AutoModelForCausalLM.from_pretrained(model_name).save_pretrained('./models/dialogpt')
+            AutoTokenizer.from_pretrained(model_name).save_pretrained('./models/dialogpt')
+
+        self.model = AutoModelForCausalLM.from_pretrained('./models/dialogpt')
+        self.tokenizer = AutoTokenizer.from_pretrained('./models/dialogpt')
+
+    def __call__(self, inputs: str) -> str:
+        inputs_tokenized = self.tokenizer.encode(inputs+ self.tokenizer.eos_token, return_tensors='pt')
+        reply_ids = self.model.generate(inputs_tokenized, max_length=1250, pad_token_id=self.tokenizer.eos_token_id)
+        reply = self.tokenizer.decode(reply_ids[:, inputs_tokenized.shape[-1]:][0], skip_special_tokens=True)
+
+        return reply
+
+    def run(self):
+        while True:
+            user_input = input("User: ")
+            print("Bot:", self(user_input))