Spaces:
Sleeping
Sleeping
File size: 1,185 Bytes
c5529b3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 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)) |