Update output format

app.py

@@ -262,24 +262,45 @@ def text_to_text_generation(verbose, prompt):
     response_num = 0 if "0" in response else (1 if "1" in response else 2)
 
     def create_generation_prompt(response_num, prompt, retriever):
+        md = ""
         match response_num:
             case 0:
                 prompt = f"[INST] {prompt}\n Lets think step by step. [/INST] {start_template}"
 
             case 1:
                 if retriever == "semantic_search":
+                    question = prompt
                     retrieved_results = semantic_search(prompt)
                     prompt = f"You are a helpful kubernetes professional. [INST] Use the following documentation, if it is relevant to answer the question below. [/INST]\nDocumentation: {retrieved_results} </s>\n<s> [INST] Answer the following question: {prompt} [/INST]\nAnswer: "
+
+                    md = (
+                        f"### Step 1: Preparing prompt for additional documentation\n"
+                        f"You are a helpful kubernetes professional. [INST] Use the following documentation, if it is relevant to answer the question below. [/INST]\nDocumentation:\n"
+                        f"### Step 2: Retrieving documentation from a book.\n"
+                        f"{retrieved_results}\n"
+                        f"### Step 3: Creating full prompt given to model\n"
+                        f"You are a helpful kubernetes professional. [INST] Use the following documentation, if it is relevant to answer the question below. [/INST]\nDocumentation: [RETRIEVED_RESULTS_FROM_BOOK] [INST] Answer the following question: {question} [/INST]\nAnswer:  \n"
+                    )
                 elif retriever == "google_search":
                     retrieved_results = google_search(prompt)
+                    question = prompt
                     prompt = f"You are a helpful kubernetes professional. [INST] Use the following documentation, if it is relevant to answer the question below. [/INST]\nDocumentation: {retrieved_results} </s>\n<s> [INST] Answer the following question: {prompt} [/INST]\nAnswer: "
+
+                    md = (
+                        f"### Step 1: Preparing prompt for additional documentation\n"
+                        f"You are a helpful kubernetes professional. [INST] Use the following documentation, if it is relevant to answer the question below. [/INST]\nDocumentation:\n"
+                        f"### Step 2: Retrieving documentation from Google.\n"
+                        f"{retrieved_results}\n"
+                        f"### Step 3: Creating full prompt given to model\n"
+                        f"You are a helpful kubernetes professional. [INST] Use the following documentation, if it is relevant to answer the question below. [/INST]\nDocumentation: [RETRIEVED_RESULTS_FROM_GOOGLE] [INST] Answer the following question: {question} [/INST]\nAnswer:  \n"
+                    )
                 else:
                     prompt = f"[INST] Answer the following question: {prompt} [/INST]\nAnswer: "
 
             case _:
                 prompt = f"[INST] {prompt} [/INST]"
 
-        return prompt
+        return prompt, md
 
     def generate_batch(*prompts):
         tokenized_inputs = tokenizer(prompts, return_tensors="pt", padding=True).to("cuda")
@@ -307,7 +328,7 @@ def text_to_text_generation(verbose, prompt):
 
         return (prompt, answer[start:end].strip())
 
-    modes = ["Kubectl command", "Kubernetes definition", "Normal"]
+    modes = ["Kubectl command", "Kubernetes definition", "Other"]
 
     print(f'{" Query Start ":-^40}')
     print("Classified as: " + modes[response_num])
@@ -317,23 +338,29 @@ def text_to_text_generation(verbose, prompt):
 
 
     if response_num == 2:
-        prompt = create_generation_prompt(response_num, prompt, False)
+        prompt, md = create_generation_prompt(response_num, prompt, False)
         print('Prompt given to model:\n' + prompt + '\n')
         original, new = generate_batch(prompt)[0]
         prompt, response = cleanup(original, new)
         if verbose:
             return (
                 f"# 📚KubeWizard📚\n"
-                f"{modes}\n\n" f"# Prompt given to the model:\n" f"{str_to_md(prompt)}\n" f"# Model's answer:\n" f"{str_to_md(response)}\n"
+                f"# Classified your prompt as:\n"
+                f"{modes}\n\n" 
+                f"# Prompt given to the model:\n" 
+                f"{str_to_md(prompt)}\n"
+                f"# Model's answer:\n" f"{str_to_md(response)}\n"
             )
         else:
             return (
                 f"# 📚KubeWizard📚\n"
-                f"{modes}\n\n" f"# Answer:\n" f"{str_to_md(response)}"
+                f"# Classified your prompt as:\n"
+                f"{modes}\n\n" 
+                f"# Answer:\n" f"{str_to_md(response)}"
             )
 
     if response_num == 0:
-        prompt = create_generation_prompt(response_num, prompt, False)
+        prompt, md = create_generation_prompt(response_num, prompt, False)
         print('Prompt given to model:\n' + prompt + '\n')
         original, new = generate_batch(prompt)[0]
         prompt, response = cleanup(original, new)
@@ -341,6 +368,7 @@ def text_to_text_generation(verbose, prompt):
         if verbose:
             return (
                 f"# 📚KubeWizard📚\n"
+                f"# Classified your prompt as:\n"
                 f"{modes}\n\n"
                 f"# Prompt given to the model:\n"
                 f"{str_to_md(prompt)}\n"
@@ -352,14 +380,16 @@ def text_to_text_generation(verbose, prompt):
         else:
             return (
                 f"# 📚KubeWizard📚\n"
-                f"{modes}\n\n" f"# Answer:\n" f"```bash\n{str_to_md(response)}\n```\n"
+                f"# Classified your prompt as:\n"
+                f"{modes}\n\n"
+                f"# Answer:\n" f"```bash\n{str_to_md(response)}\n```\n"
             )
 
-    res_prompt = create_generation_prompt(response_num, prompt, False)
+    res_prompt, res_md = create_generation_prompt(response_num, prompt, False)
     print(f'Prompt given to finetuned model:\n{res_prompt}\n')
-    res_semantic_search_prompt = create_generation_prompt(response_num, prompt, "semantic_search")
+    res_semantic_search_prompt, res_semantic_search_md = create_generation_prompt(response_num, prompt, "semantic_search")
     print(f'Prompt given to model with RAG:\n{res_semantic_search_prompt}\n')
-    res_google_search_prompt = create_generation_prompt(response_num, prompt, "google_search")
+    res_google_search_prompt, res_google_search_md = create_generation_prompt(response_num, prompt, "google_search")
     print(f'Prompt given to model with Google search:\n{res_google_search_prompt}\n')
 
     gen_normal, gen_semantic_search, gen_google_search = generate_batch(
@@ -373,26 +403,31 @@ def text_to_text_generation(verbose, prompt):
     if verbose:
         return (
             f"# 📚KubeWizard📚\n"
+            f"# Classified your prompt as:\n"
             f"{modes}\n\n"
+            f"--------------------------------------------"
             f"# Answer with finetuned model\n"
             f"## Prompt given to the model:\n"
             f"{str_to_md(res_prompt)}\n\n"
             f"## Model's answer:\n"
             f"{str_to_md(res_normal)}\n\n"
+            f"--------------------------------------------"
             f"# Answer with RAG\n"
-            f"## Prompt given to the model:\n"
-            f"{str_to_md(res_semantic_search_prompt)}\n\n"
-            f"## Model's answer:\n"
-            f"{str_to_md(res_semantic_search)}\n\n"
+            f"## Section 1: Preparing for generation\n"
+            f"{res_semantic_search_md}\n"
+            f"## Section 2: Generating answer\n"
+            f"{str_to_md(res_semantic_search)}\n"
+            f"--------------------------------------------"
             f"# Answer with Google search\n"
-            f"## Prompt given to the model:\n"
-            f"{str_to_md(res_google_search_prompt)}\n\n"
-            f"## Model's answer:\n"
-            f"{str_to_md(res_google_search)}\n\n"
+            f"## Section 1: Preparing for generation\n"
+            f"{res_google_search_md}\n"
+            f"## Section 2: Generating answer\n"
+            f"{str_to_md(res_google_search)}\n"
         )
     else:
         return (
             f"# 📚KubeWizard📚\n"
+            f"# Classified your prompt as:\n"
             f"{modes}\n\n"
             f"# Answer with finetuned model\n\n {str_to_md(res_normal)}\n"
             f"# Answer with RAG\n\n {str_to_md(res_semantic_search)}\n"