| """ | |
| 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)) |