added prompts

assets/prompts/custom_prompts.py

@@ -0,0 +1,20 @@
+text_qa_template_str = (
+    "You are a helpful assistant that can give users information about california water regulations."
+    "Your task is to answer the question based on the context below and using your own knowledge.\n"
+    "Context: {context_str}\n"
+    "Question: {query_str}\n"
+    "Answer: \n"
+    "If the question is relevant, provide the name of the chapter, the article and the section."
+    "If possible, also provide the page number of the document from which your answer was found.\n"
+)
+
+refine_template_str = (
+    "The original question is as follows: {query_str}\n"
+    "We have provided an existing answer: {existing_answer}\n"
+    "We have the opportunity to refine the existing answer "
+    "(only if needed) with some more context below.\n"
+    "------------\n"
+    "{context_msg}\n"
+    "------------\n"
+    "Using both the new context and your own knowledege, update or repeat the existing answer.\n"
+)

models/llamaCustom.py

@@ -21,8 +21,11 @@ from llama_index import (
     load_index_from_storage,
 )
 from llama_index.llms import CompletionResponse, CustomLLM, LLMMetadata
+from llama_index.prompts import Prompt
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
+from assets.prompts import custom_prompts
+
 load_dotenv()
 # openai.api_key = os.getenv("OPENAI_API_KEY")
 fs = HfFileSystem()
@@ -35,6 +38,10 @@ NUM_OUTPUT = 525
 # set maximum chunk overlap
 CHUNK_OVERLAP_RATION = 0.2
 
+text_qa_template = Prompt(custom_prompts.text_qa_template_str)
+
+refine_template = Prompt(custom_prompts.refine_template_str)
+
 
 @st.cache_resource
 def load_model(model_name: str):
@@ -136,7 +143,10 @@ class LlamaCustom:
 
     def get_response(self, query_str):
         print("query_str: ", query_str)
-        query_engine = self.vector_index.as_query_engine()
+        # query_engine = self.vector_index.as_query_engine()
+        query_engine = self.vector_index.as_query_engine(
+            text_qa_template=text_qa_template, refine_template=refine_template
+        )
         response = query_engine.query(query_str)
         print("metadata: ", response.metadata)
         return str(response)