added comments
Browse files- chatbot.py +16 -6
chatbot.py
CHANGED
@@ -4,48 +4,58 @@ from pinecone import Pinecone
|
|
4 |
|
5 |
device = 'cpu'
|
6 |
|
|
|
7 |
pc = Pinecone(api_key='89eeb534-da10-4068-92f7-12eddeabe1e5')
|
8 |
|
|
|
9 |
index_name = 'abstractive-question-answering'
|
10 |
index = pc.Index(index_name)
|
11 |
|
|
|
12 |
def load_models():
|
13 |
-
print("Loading models...")
|
14 |
-
|
15 |
retriever = SentenceTransformer("flax-sentence-embeddings/all_datasets_v3_mpnet-base")
|
16 |
tokenizer = T5Tokenizer.from_pretrained('t5-small')
|
17 |
generator = T5ForConditionalGeneration.from_pretrained('t5-base').to(device)
|
18 |
-
|
19 |
return retriever, generator, tokenizer
|
20 |
|
21 |
retriever, generator, tokenizer = load_models()
|
22 |
|
23 |
def process_query(query):
|
|
|
24 |
xq = retriever.encode([query]).tolist()
|
|
|
25 |
xc = index.query(vector=xq, top_k=1, include_metadata=True)
|
26 |
|
27 |
print("Pinecone response:", xc)
|
28 |
|
|
|
29 |
if 'matches' in xc and isinstance(xc['matches'], list):
|
30 |
context = [m['metadata']['Output'] for m in xc['matches']]
|
31 |
context_str = " ".join(context)
|
32 |
formatted_query = f"answer the question: {query} context: {context_str}"
|
33 |
|
34 |
-
#
|
35 |
output_text = context_str
|
36 |
if len(output_text.splitlines()) > 5:
|
37 |
return output_text
|
38 |
|
|
|
39 |
if output_text.lower() == "none":
|
40 |
return "The topic is not covered in the student manual."
|
41 |
|
|
|
42 |
inputs = tokenizer.encode(formatted_query, return_tensors="pt", max_length=512, truncation=True).to(device)
|
|
|
43 |
ids = generator.generate(inputs, num_beams=2, min_length=10, max_length=60, repetition_penalty=1.2)
|
|
|
44 |
answer = tokenizer.decode(ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
45 |
|
46 |
-
|
47 |
-
|
|
|
48 |
if any(keyword in answer.lower() for keyword in nli_keywords):
|
49 |
return context_str
|
50 |
|
|
|
51 |
return answer
|
|
|
4 |
|
5 |
device = 'cpu'
|
6 |
|
7 |
+
# Calling the pinecone api
|
8 |
pc = Pinecone(api_key='89eeb534-da10-4068-92f7-12eddeabe1e5')
|
9 |
|
10 |
+
# Connect to the Pinecone index for querying and storing vectors
|
11 |
index_name = 'abstractive-question-answering'
|
12 |
index = pc.Index(index_name)
|
13 |
|
14 |
+
# Load the retriever model for sentence embeddings and the T5 model for text generation
|
15 |
def load_models():
|
16 |
+
print("Loading models...")
|
|
|
17 |
retriever = SentenceTransformer("flax-sentence-embeddings/all_datasets_v3_mpnet-base")
|
18 |
tokenizer = T5Tokenizer.from_pretrained('t5-small')
|
19 |
generator = T5ForConditionalGeneration.from_pretrained('t5-base').to(device)
|
|
|
20 |
return retriever, generator, tokenizer
|
21 |
|
22 |
retriever, generator, tokenizer = load_models()
|
23 |
|
24 |
def process_query(query):
|
25 |
+
# Encode the query into a vector for semantic search using SentenceTransformer
|
26 |
xq = retriever.encode([query]).tolist()
|
27 |
+
# Query the Pinecone index for the most similar vector to the query
|
28 |
xc = index.query(vector=xq, top_k=1, include_metadata=True)
|
29 |
|
30 |
print("Pinecone response:", xc)
|
31 |
|
32 |
+
# Concatenates the original question with the context extracted from the matched metadata
|
33 |
if 'matches' in xc and isinstance(xc['matches'], list):
|
34 |
context = [m['metadata']['Output'] for m in xc['matches']]
|
35 |
context_str = " ".join(context)
|
36 |
formatted_query = f"answer the question: {query} context: {context_str}"
|
37 |
|
38 |
+
# If the context is longer than 5 lines, return the context extracted from Pinecone directly
|
39 |
output_text = context_str
|
40 |
if len(output_text.splitlines()) > 5:
|
41 |
return output_text
|
42 |
|
43 |
+
# If none, then it will return that it was not covered in the student manual
|
44 |
if output_text.lower() == "none":
|
45 |
return "The topic is not covered in the student manual."
|
46 |
|
47 |
+
# Tokenizes the formatted query
|
48 |
inputs = tokenizer.encode(formatted_query, return_tensors="pt", max_length=512, truncation=True).to(device)
|
49 |
+
# Generates an answer using the t5 model
|
50 |
ids = generator.generate(inputs, num_beams=2, min_length=10, max_length=60, repetition_penalty=1.2)
|
51 |
+
# Decodes the answer to make it readable for the user
|
52 |
answer = tokenizer.decode(ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
|
53 |
|
54 |
+
|
55 |
+
# If it has this words, it will just paste the output from the extracted meta-data output from pinecone
|
56 |
+
nli_keywords = ['not_equivalent', 'not_entailment', 'entailment', 'neutral', 'not_enquiry']
|
57 |
if any(keyword in answer.lower() for keyword in nli_keywords):
|
58 |
return context_str
|
59 |
|
60 |
+
# returns the answer
|
61 |
return answer
|