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