from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModel import json import spacy from celebbot import CelebBot from utils import * DEBUG = False 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()) if not DEBUG: ai.text_to_speech() ai.text = "" # Run the AI if __name__ == "__main__": main()