Adding QueryModification Pipeline

Query_Modification/QueryModification.py

@@ -0,0 +1,121 @@
+import json
+import pandas as pd
+import google.generativeai as genai
+
+# Function to process text input with Gemini model
+def query_Modifier(input_text):
+
+    with open('config.json', 'r') as file:
+        config = json.load(file)
+
+    gemini_key = config.get("GEMINI")
+
+    # Initialize the API key
+    genai.configure(api_key=gemini_key)
+
+    # print(gemini_key)
+
+    # Load the prompt from file
+    with open("Query_Modification/prompt.txt", 'r') as file:
+        PROMPT_TEMPLATE = file.read()
+
+    # Safety settings for Gemini model
+    safe = [
+        {
+            "category": "HARM_CATEGORY_DANGEROUS",
+            "threshold": "BLOCK_NONE",
+        },
+        {
+            "category": "HARM_CATEGORY_HARASSMENT",
+            "threshold": "BLOCK_NONE",
+        },
+        {
+            "category": "HARM_CATEGORY_HATE_SPEECH",
+            "threshold": "BLOCK_NONE",
+        },
+        {
+            "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+            "threshold": "BLOCK_NONE",
+        },
+        {
+            "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+            "threshold": "BLOCK_NONE",
+        },
+    ]
+
+    generation_config = {
+        "temperature": 1,
+        "top_p": 0.95,
+        "top_k": 40,
+        "max_output_tokens": 8192,
+        "response_mime_type": "text/plain",
+    }
+    
+    # Initialize the generative model
+    model = genai.GenerativeModel("gemini-1.5-flash", generation_config=generation_config)
+
+    
+    full_prompt = f"{input_text}\n\n{PROMPT_TEMPLATE}"
+                
+    # Call the generative model for text input
+    result = model.generate_content([full_prompt], safety_settings=safe)
+    return result.text
+
+
+def getKeywords(input_text):
+    # Extract keywords from the input text
+
+    with open('config.json', 'r') as file:
+        config = json.load(file)
+
+    gemini_key = config.get("GEMINI")
+
+    # Initialize the API key
+    genai.configure(api_key=gemini_key)
+
+    # Safety settings for Gemini model
+    safe = [
+        {
+            "category": "HARM_CATEGORY_DANGEROUS",
+            "threshold": "BLOCK_NONE",
+        },
+        {
+            "category": "HARM_CATEGORY_HARASSMENT",
+            "threshold": "BLOCK_NONE",
+        },
+        {
+            "category": "HARM_CATEGORY_HATE_SPEECH",
+            "threshold": "BLOCK_NONE",
+        },
+        {
+            "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+            "threshold": "BLOCK_NONE",
+        },
+        {
+            "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+            "threshold": "BLOCK_NONE",
+        },
+    ]
+
+    generation_config = {
+        "temperature": 1,
+        "top_p": 0.95,
+        "top_k": 40,
+        "max_output_tokens": 8192,
+        "response_mime_type": "text/plain",
+    }
+    
+    # Initialize the generative model
+    model = genai.GenerativeModel("gemini-1.5-flash", generation_config=generation_config)
+
+    
+    full_prompt = f"{input_text} \n\n Give the Keywords for the above sentence and output nothing else."
+                
+    # Call the generative model for text input
+    result = model.generate_content([full_prompt], safety_settings=safe)
+
+    response = result.text
+    response = response.replace("Keywords:", "")
+    response = response.replace(",", "")
+
+    return response.strip()

Query_Modification/prompt.txt

@@ -0,0 +1,19 @@
+Modify the following query to improve its suitability for a Retrieval Augmented Generation (RAG) system using a semantic search engine like Cosign:
+
+Original Query: [Original query here]
+
+Guidelines:
+
+Clarity and Specificity: Make the query more specific and focused.
+Keyword Optimization: Identify and include relevant keywords that align with the dataset.
+Semantic Relevance: Consider the underlying meaning and context of the query.
+Question Formulation: Frame the query as a question to facilitate direct answer extraction.
+Contextual Clues: If applicable, provide additional context or background information.
+
+Example:
+
+Original Query: "Tell me about the French Revolution"
+
+Modified Query: "What were the main causes and effects of the French Revolution, and who were its key figures?"
+
+Guardrail : Output only the Modified Query.