adding number of chunks and context

app.py

@@ -37,13 +37,14 @@ def generate(prompt, model_parameters):
  )
 
 
-def submit_input(input_, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature):
+def submit_input(input_, num_chunks, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature):
  if input_.strip() == "":
  gr.Warning("Not possible to inference an empty input")
  return None
 
 
  model_parameters = {
+ "NUM_CHUNKS": num_chunks,
  "MAX_NEW_TOKENS": max_new_tokens,
  "REPETITION_PENALTY": repetition_penalty,
  "TOP_K": top_k,
@@ -109,6 +110,13 @@ def gradio_app():
 
  with gr.Row(variant="panel"):
  with gr.Accordion("Model parameters", open=False, visible=SHOW_MODEL_PARAMETERS_IN_UI):
+ num_chunks = Slider(
+ minimum=1,
+ maximum=6,
+ step=1,
+ value=4,
+ label="Number of chunks"
+ )
  max_new_tokens = Slider(
  minimum=50,
  maximum=1000,
@@ -154,7 +162,7 @@ def gradio_app():
  label="Temperature"
  )
 
- parameters_compontents = [max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature]
+ parameters_compontents = [num_chunks, max_new_tokens, repetition_penalty, top_k, top_p, do_sample, num_beams, temperature]
 
  with gr.Column(variant="panel"):
  output = Textbox(

rag.py

@@ -24,19 +24,11 @@ class RAG:
 
  logging.info("RAG loaded!")
  
- def get_context(self, instruction, number_of_contexts=3):
-
- context = ""
-
+ def get_context(self, instruction, number_of_contexts=4):
 
  documentos = self.vectore_store.similarity_search_with_score(instruction, k=number_of_contexts)
 
-
- for doc in documentos:
-
- context += doc[0].page_content
-
- return context
+ return documentos
  
  def predict(self, instruction, context, model_parameters):
 
@@ -61,14 +53,30 @@ class RAG:
  response = requests.post(self.model_name, headers=headers, json=payload)
 
  return response.json()[0]["generated_text"].split("###")[-1][8:-1]
+ 
+ def beautiful_context(self, docs):
+
+ text_context = ""
+
+ full_context = ""
+
+ for doc in docs:
+ text_context += doc[0].page_content
+ full_context += doc[0].page_content + "\n"
+ full_context += doc[0].metadata["Títol de la norma"] + "\n\n"
+
+ return text_context, full_context
 
  def get_response(self, prompt: str, model_parameters: dict) -> str:
  
- context = self.get_context(prompt)
+ docs = self.get_context(prompt, model_parameters["NUM_CHUNKS"])
+ text_context, full_context = beautiful_context(docs)
+
+ del model_parameters["NUM_CHUNKS"]
 
- response = self.predict(prompt, context, model_parameters)
+ response = self.predict(prompt, text_context, model_parameters)
 
  if not response:
  return self.NO_ANSWER_MESSAGE
 
- return response
+ return response, full_context