Spaces:
Runtime error
Runtime error
File size: 1,991 Bytes
6bc94ac abca9bf 6bc94ac 5beab45 6bc94ac 5beab45 6bc94ac 5beab45 6bc94ac abca9bf 6bc94ac |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModel
import json
import spacy
from celebbot import CelebBot
from utils import *
DEBUG = True
QA_MODEL_ID = "google/flan-t5-large"
SENTTR_MODEL_ID = "sentence-transformers/all-mpnet-base-v2"
def main():
celeb_data = get_celeb_data("data.json")
message = "Please choose your favorite celebrity from\n"\
"1. Cate Blanchett\n"\
"2. David Beckham\n"\
"3. Emma Watson\n"\
"4. Lady Gaga\n"\
"5. Madonna\n"\
"6. Mark Zuckerberg\n"\
"input name:\n"
QA_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL_ID)
QA_model = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL_ID)
sentTr_tokenizer = AutoTokenizer.from_pretrained(SENTTR_MODEL_ID)
sentTr_model = AutoModel.from_pretrained(SENTTR_MODEL_ID)
celeb_name = input(message)
gender = celeb_data[celeb_name]["gender"]
if celeb_name == "Madonna":
name = "Madonna-American-singer-and-actress"
elif celeb_name == "Anne Hathaway":
name = "Anne-Hathaway-American-actress"
else:
name="-".join(celeb_name.split(" "))
knowledge = get_article(f"https://www.britannica.com/biography/{name}")
spacy_model = spacy.load("en_core_web_lg")
knowledge_sents = [i.text.strip() for i in spacy_model(knowledge).sents]
ai = CelebBot(celeb_name, gender, QA_tokenizer, QA_model, sentTr_tokenizer, sentTr_model, spacy_model, knowledge_sents)
answers = []
while True:
if not DEBUG:
ai.speech_to_text()
else:
ai.text = input("Your question: ")
pass
# ai.text = q
if ai.text != "":
print("me --> ", ai.text)
answers.append(ai.question_answer())
print("bot --> ", ai.text)
if not DEBUG:
ai.text_to_speech()
ai.text = ""
# Run the AI
if __name__ == "__main__":
main() |