Update semantic search and output format

app.py

@@ -72,9 +72,9 @@ def get_results_from_pinecone(query, top_k=3, re_rank=True, verbose=True):
     return final_results
 
 def semantic_search(prompt):
-    final_results = get_results_from_pinecone(prompt, top_k=3, re_rank=True, verbose=True)
+    final_results = get_results_from_pinecone(prompt, top_k=9, re_rank=True, verbose=True)
 
-    return 'First result:\n' + final_results[0]['metadata']['text'].replace('\n', ' ') + '\n' + 'Second result:\n' + final_results[1]['metadata']['text'].replace('\n', ' ') + '\n' + 'Third result:\n' + final_results[2]['metadata']['text'].replace('\n', ' ')
+    return '\n\n'.join(res['metadata']['text'].strip() for res in final_results[:3])
 
 cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
 sentencetransformer_model = SentenceTransformer('sentence-transformers/multi-qa-mpnet-base-cos-v1')
@@ -129,7 +129,7 @@ stop_terms=["</s>", "#End"]
 eos_token_ids_custom =  [torch.tensor(tokenizer.encode(term, add_special_tokens=False)).to("cuda") for term in stop_terms]
 
 category_terms=["</s>", "\n"]
-category_eos_token_ids_custom =  [torch.tensor(tokenizer.encode(term, add_special_tokens=False)).to("cuda") for term in stop_terms]
+category_eos_token_ids_custom =  [torch.tensor(tokenizer.encode(term, add_special_tokens=False)).to("cuda") for term in category_terms]
 
 
 class EvalStopCriterion(StoppingCriteria):
@@ -184,7 +184,7 @@ def text_to_text_generation(prompt):
         print(f'[INST] You are an assistant who summarizes results retrieved from a book about Kubernetes. This summary should answer the question. If the answer is not in the retrieved results, use your general knowledge. [/INST] Question: {prompt}\nRetrieved results:\n{retrieved_results}\nResponse:')
         prompt = f'[INST] You are an assistant who summarizes results retrieved from a book about Kubernetes. This summary should answer the question. If the answer is not in the retrieved results, use your general knowledge. [/INST] Question: {prompt}\nRetrieved results:\n{retrieved_results}\nResponse:'
     else:
-        prompt = f'[INST] {prompt} [/INST]'
+        prompt = f'[INST] {prompt} [/INST]'
 
     # Generate output
     model_input = tokenizer(prompt, return_tensors="pt").to("cuda")