Update output format, batch inference

app.py

@@ -1,3 +1,4 @@
+import re
 import torch
 import time
 import pinecone
@@ -10,36 +11,47 @@ from transformers.generation.stopping_criteria import StoppingCriteria, Stopping
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from torch import nn
 from sentence_transformers.cross_encoder import CrossEncoder
-from peft import PeftModel
 from sentence_transformers import SentenceTransformer
+from peft import PeftModel
 from bs4 import BeautifulSoup
 import requests
 
 headers = {
-  "User-Agent":"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_5) AppleWebKit 537.36 (KHTML, like Gecko) Chrome",
-  "Accept":"text/html,application/xhtml+xml,application/xml; q=0.9,image/webp,*/*;q=0.8",
-  'Cookie':'CONSENT=YES+cb.20210418-17-p0.it+FX+917; '
+    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_5) AppleWebKit 537.36 (KHTML, like Gecko) Chrome",
+    "Accept": "text/html,application/xhtml+xml,application/xml; q=0.9,image/webp,*/*;q=0.8",
+    "Cookie": "CONSENT=YES+cb.20210418-17-p0.it+FX+917; ",
 }
 
+
 def google_search(text):
     print(f"Google search on: {text}")
     try:
-        site = requests.get(f'https://www.google.com/search?hl=en&q={text}', headers=headers)
-        main = BeautifulSoup(site.text, features="html.parser").select_one('#main').select('.VwiC3b.lyLwlc.yDYNvb.W8l4ac')
-        res =  '\n\n'.join([m.get_text() for m in main])
+        site = requests.get(f"https://www.google.com/search?hl=en&q={text}", headers=headers)
+        main = (
+            BeautifulSoup(site.text, features="html.parser").select_one("#main").select(".VwiC3b.lyLwlc.yDYNvb.W8l4ac")
+        )
+        res = []
+        for m in main:
+            t = m.get_text()
+            if "—" in t:
+                t = t[len("—") + t.index("—") :].strip()
+
+            res.append(t)
+
+        ans = "\n".join(res)
     except Exception as ex:
         print(f"Error: {ex}")
-        res = ""
+        ans = ""
 
-    print(f"The result of the google search is: {res}")
+    print(f"The result of the google search is: {ans}")
 
-    return res
+    return ans
 
 PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
 
-sentencetransformer_model = SentenceTransformer('sentence-transformers/multi-qa-mpnet-base-cos-v1')
 pinecone.init(api_key=PINECONE_API_KEY, environment="gcp-starter")
 
+sentencetransformer_model = SentenceTransformer('sentence-transformers/multi-qa-mpnet-base-cos-v1')
 
 CACHE_DIR = "./.cache"
 INDEX_NAME = "k8s-semantic-search"
@@ -79,6 +91,7 @@ def create_embedding(text: str):
     
     return embed_text.tolist()
 
+
 index = pinecone.Index(INDEX_NAME)
 
 
@@ -190,6 +203,23 @@ start_template = "### Answer:"
 command_template = "# Command:"
 end_template = "#End"
 
+def str_to_md(text):
+    def escape_hash(line):
+        i = 0
+        while i < len(line) and line[i] == ' ':
+            i+=1
+
+        if i == len(line):
+            return line
+        
+        if line[i] == '#':
+            line = line[:i] + '\\' + line[i:]
+        
+        return line
+
+    lines = text.split('\n')
+    lines = [escape_hash(line) for line in lines]
+    return '  \n'.join(lines)
 
 def text_to_text_generation(verbose, prompt):
     prompt = prompt.strip()
@@ -205,6 +235,7 @@ def text_to_text_generation(verbose, prompt):
     )
 
     model_input = tokenizer(is_kubectl_prompt, return_tensors="pt").to("cuda")
+
     with torch.no_grad():
         response = tokenizer.decode(
             model.generate(
@@ -223,7 +254,7 @@ def text_to_text_generation(verbose, prompt):
 
     response_num = 0 if "0" in response else (1 if "1" in response else 2)
 
-    def generate(response_num, prompt, retriever, verbose):
+    def create_generation_prompt(response_num, prompt, retriever):
         match response_num:
             case 0:
                 prompt = f"[INST] {prompt}\n Lets think step by step. [/INST] {start_template}"
@@ -241,56 +272,104 @@ def text_to_text_generation(verbose, prompt):
             case _:
                 prompt = f"[INST] {prompt} [/INST]"
 
-        print("Query:")
-        print(prompt)
+        return prompt
+
+    def generate_batch(*prompts):
+        tokenized_inputs = tokenizer(prompts, return_tensors="pt", padding=True).to("cuda")
 
-        # Generate output
-        model_input = tokenizer(prompt, return_tensors="pt").to("cuda")
         with torch.no_grad():
-            response = tokenizer.decode(
+            responses = tokenizer.batch_decode(
                 model.generate(
-                    **model_input,
+                    **tokenized_inputs,
                     max_new_tokens=256,
                     pad_token_id=tokenizer.eos_token_id,
                     repetition_penalty=1.15,
                     stopping_criteria=StoppingCriteriaList([eval_stop_criterion]),
-                )[0],
+                ),
                 skip_special_tokens=True,
             )
 
-        decoded_prompt = tokenizer.decode(tokenizer(prompt).input_ids, skip_special_tokens=True)
+        decoded_prompts = tokenizer.batch_decode(tokenized_inputs.input_ids, skip_special_tokens=True)
 
-        start = (
-            response.index(start_template) + len(start_template) if start_template in response else len(decoded_prompt)
-        )
-        start = response.index(command_template) + len(command_template) if command_template in response else start
-        end = response.index(end_template) if end_template in response else len(response)
+        return [(prompt, answer) for prompt, answer in zip(decoded_prompts, responses)]
 
-        return response if verbose else response[start:end].strip()
+    def cleanup(prompt, answer):
+        start = answer.index(start_template) + len(start_template) if start_template in answer else len(prompt)
+        start = answer.index(command_template) + len(command_template) if command_template in answer else start
+        end = answer.index(end_template) if end_template in answer else len(answer)
 
-    true_response = generate(response_num, prompt, False, verbose)
-    true_response_semantic_search = generate(response_num, prompt, "semantic_search", verbose)
-    true_response_google_search = generate(response_num, prompt, "google_search", verbose)
+        return (prompt, answer[start:end].strip())
 
+    modes = ["Kubectl command", "Kubernetes definition", "Normal"]
+    modes[response_num] = f"**{modes[response_num]}**"
+    modes = " / ".join(modes)
 
-    print("Returned: " + true_response)
-    print(f'{" QUERY END ":-^40}')
+    if response_num == 2:
+        prompt = create_generation_prompt(response_num, prompt, False)
+        original, new = generate_batch(prompt)[0]
+        prompt, response = cleanup(original, new)
+        if verbose:
+            return 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"{modes}\n\n" f"# Answer:\n" f"{str_to_md(response)}"
+
+    if response_num == 0:
+        prompt = create_generation_prompt(response_num, prompt, False)
+        original, new = generate_batch(prompt)[0]
+        prompt, response = cleanup(original, new)
+        model_response = new[len(original):].strip()
+        if verbose:
+            return (
+                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(model_response)}\n"
+                f"# Processed answer:\n"
+                f"```bash\n{str_to_md(response)}\n```\n"
+            )
+        else:
+            return f"{modes}\n\n" f"# Answer:\n" f"```bash\n{str_to_md(response)}\n```\n"
 
-    match response_num:
-        case 0:
-            mode = "Kubectl"
-        case 1:
-            mode = "Kubernetes"
-        case _:
-            mode = "Normal"
+    res_prompt = create_generation_prompt(response_num, prompt, False)
+    res_semantic_search_prompt = create_generation_prompt(response_num, prompt, "semantic_search")
+    res_google_search_prompt = create_generation_prompt(response_num, prompt, "google_search")
 
-    return (
-        f"*Mode*: {mode}",
-        f"# Answer\n\n {true_response}",
-        f"# Answer with RAG\n\n {true_response_semantic_search}",
-        f"# Answer with Google search\n\n {true_response_google_search}"
+    gen_normal, gen_semantic_search, gen_google_search = generate_batch(
+        res_prompt, res_semantic_search_prompt, res_google_search_prompt
     )
 
+    res_prompt, res_normal = cleanup(*gen_normal)
+    res_semantic_search_prompt, res_semantic_search = cleanup(*gen_semantic_search)
+    res_google_search_prompt, res_google_search = cleanup(*gen_google_search)
+
+    if verbose:
+        return (
+            f"{modes}\n\n"
+            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"# 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"# 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"
+        )
+    else:
+        return (
+            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"
+            f"# Answer with Google search\n\n {str_to_md(res_google_search)}"
+        )
+
 
 iface = gr.Interface(
     fn=text_to_text_generation,
@@ -298,12 +377,7 @@ iface = gr.Interface(
         gr.components.Checkbox(label="Verbose"),
         gr.components.Text(placeholder="prompt here ...", label="Prompt"),
     ],
-    outputs=[
-        gr.components.Markdown(label="Mode"),
-        gr.components.Markdown(label="Answer Without Retriever"),
-        gr.components.Markdown(label="Answer With Retriever"),
-        gr.components.Markdown(label="Answer With Google search"),
-    ],
+    outputs=gr.components.Markdown(label="Answer"),
     allow_flagging="never",
 )