Spaces:
Runtime error
Runtime error
File size: 2,747 Bytes
706771c |
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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
from collections import deque
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from sentence_transformers import SentenceTransformer
from utils import generate_response
import pandas as pd
import pickle
from utils import encode_rag, cosine_sim_rag, top_candidates
class ChatBot:
def __init__(self):
self.conversation_history = deque([], maxlen=10)
self.generative_model = None
self.generative_tokenizer = None
self.vect_data = []
self.scripts = []
self.ranking_model = None
def load(self):
""" "This method is called first to load all datasets and
model used by the chat bot; all the data to be saved in
tha data folder, models to be loaded from hugging face"""
with open("data/scripts_vectors.pkl", "rb") as fp:
self.vect_data = pickle.load(fp)
self.scripts = pd.read_pickle("data/scripts.pkl")
self.ranking_model = SentenceTransformer(
"Shakhovak/chatbot_sentence-transformer"
)
self.generative_model = AutoModelForSeq2SeqLM.from_pretrained(
"Shakhovak/flan-t5-base-sheldon-chat-v2"
)
self.generative_tokenizer = AutoTokenizer.from_pretrained(
"Shakhovak/flan-t5-base-sheldon-chat-v2"
)
def generate_response(self, utterance):
query_encoding = encode_rag(
texts=utterance,
model=self.ranking_model,
contexts=self.conversation_history,
)
bot_cosine_scores = cosine_sim_rag(
self.vect_data,
query_encoding,
)
top_scores, top_indexes = top_candidates(
bot_cosine_scores, initial_data=self.scripts
)
if top_scores[0] >= 0.89:
for index in top_indexes:
rag_answer = self.scripts.iloc[index]["answer"]
answer = generate_response(
model=self.generative_model,
tokenizer=self.generative_tokenizer,
question=utterance,
context=self.conversation_history,
top_p=0.9,
temperature=0.95,
rag_answer=rag_answer,
)
else:
answer = generate_response(
model=self.generative_model,
tokenizer=self.generative_tokenizer,
question=utterance,
context=self.conversation_history,
top_p=0.9,
temperature=0.95,
)
self.conversation_history.append(utterance)
self.conversation_history.append(answer)
return answer
# katya = ChatBot()
# katya.load()
# print(katya.generate_response("What is he doing there?"))
|