Create models/blenderbot.py

models/blenderbot.py

@@ -0,0 +1,33 @@
+"""
+Adapted from:
+https://mandgie.medium.com/how-to-build-your-own-chatbot-f5848ebcba8d
+"""
+
+
+from transformers import BlenderbotSmallTokenizer, BlenderbotSmallForConditionalGeneration
+import os
+
+
+class BlenderBot:
+    def __init__(
+        self,
+        model_name: str ='facebook/blenderbot_small-90M',
+    ):
+        if not os.path.exists('./models/blenderbot'):
+            BlenderbotSmallForConditionalGeneration.from_pretrained(model_name).save_pretrained('./models/blenderbot')
+            BlenderbotSmallTokenizer.from_pretrained(model_name).save_pretrained('./models/blenderbot')
+
+        self.model = BlenderbotSmallForConditionalGeneration.from_pretrained('./models/blenderbot')
+        self.tokenizer = BlenderbotSmallTokenizer.from_pretrained('./models/blenderbot')
+
+    def __call__(self, inputs: str) -> str:
+        inputs_tokenized = self.tokenizer(inputs, return_tensors='pt')
+        reply_ids = self.model.generate(**inputs_tokenized)
+        reply = self.tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]
+
+        return reply
+
+    def run(self):
+        while True:
+            user_input = input("User: ")
+            print("Bot:", self(user_input))