Added search engine prompt engineering

hanlde_form_submit.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from grouped_sampling import GroupedSamplingPipeLine
 
+from prompt_engeneering import rewrite_prompt
 from supported_models import get_supported_model_names
 
 
@@ -15,7 +16,7 @@ def create_pipeline(model_name: str, group_size) -> GroupedSamplingPipeLine:
     :return: A pipeline with the given model name and group size.
     """
     print(f"Starts downloading model: {model_name} from the internet.")
-    pipeline =  GroupedSamplingPipeLine(
+    pipeline = GroupedSamplingPipeLine(
         model_name=model_name,
         group_size=group_size,
         end_of_sentence_stop=True,
@@ -38,9 +39,12 @@ def generate_text(
     :param output_length: The size of the groups to use.
     :return: The generated text.
     """
+    better_prompt = rewrite_prompt(prompt, 5)
     return pipeline(
-        prompt_s=prompt,
+        prompt_s=better_prompt,
         max_new_tokens=output_length,
+        return_text=True,
+        return_full_text=False,
     )["generated_text"]
 
 

prompt_engeneering.py

@@ -0,0 +1,53 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Generator
+
+import requests
+
+
+@dataclass
+class SearchResult:
+    def __init__(self, title: str, body: str, url: str):
+        self.title = title
+        self.body = body
+        self.url = url
+
+
+def get_web_search_results(
+        prompt: str,
+        num_results: int,
+) -> Generator[SearchResult, None, None]:
+    """Adds web search results to the prompt.
+    Using """
+    url = f"https://ddg-webapp-aagd.vercel.app/search?max_results={num_results}&q=${prompt}"
+    response = requests.get(url)
+    if response.status_code != 200:
+        raise ValueError(f"Failed to get web search results for prompt: {prompt}")
+    results = response.json()
+    for result in results:
+        yield SearchResult(
+            title=result["title"],
+            body=result["body"],
+            url=result["href"],
+        )
+
+
+def format_search_result(search_result: Generator[SearchResult, None, None]) -> str:
+    """Formats a search result to be added to the prompt."""
+    ans = ""
+    for i, result in enumerate(search_result):
+        ans += f"[{i}] {result.body}\nURL: {result.url}\n\n"
+    return ans
+
+
+def rewrite_prompt(
+        prompt: str,
+        num_results: int,
+) -> str:
+    """Rewrites the prompt by adding web search results to it."""
+    raw_results = get_web_search_results(prompt, num_results)
+    formatted_results = "Web search results: " + format_search_result(raw_results)
+    formatted_date = "Current date: " + datetime.now().strftime("%d/%m/%Y")
+    default_instructions = "Instructions: Using the provided web search results, write a comprehensive reply to the given query. Make sure to cite results using [[number](URL)] notation after the reference. If the provided search results refer to multiple subjects with the same name, write separate answers for each subject."
+    formatted_prompt = f"Query: {prompt}"
+    return "\n".join([formatted_results, formatted_date, default_instructions, formatted_prompt])