Update: Basic building block ready

src/app.py

@@ -0,0 +1,15 @@
+import os
+
+load_dotenv()
+import gradio as gr
+from init import Initializer
+from dotenv import load_dotenv
+
+AUG_TOKEN = os.environ.get("AUG_TOKEN")
+RES_TOKEN = os.environ.get("RES_TOKEN")
+chroma_filename = ""
+brain = Initializer.initialize(AUG_TOKEN, RES_TOKEN, chroma_filename)
+
+# TODO:
+# Chatbot like UI
+# Multiple PDF file handling ability

src/brain.py

@@ -10,48 +10,90 @@ from dotenv import load_dotenv
 
 load_dotenv()
 
+logging.basicConfig(
+    filename="bot_errors.log",
+    level=logging.ERROR,
+    format="%(asctime)s - %(levelname)s - %(message)s",
+)
+
 
 class GeminiEmbeddingFunction(EmbeddingFunction):
     def __call__(self, input: Documents) -> Embeddings:
         model = "models/embedding-001"
-        title = "Custom Queries"
+        title = "Custom query"
         return genai.embed_content(
             model=model, content=input, task_type="retrieval_document", title=title
         )["embedding"]
 
-
 class Brain:
     def __init__(
         self,
-        aug_model_name,
-        res_model_name,
-        aug_config,
-        res_config,
-        aug_model_key,
-        res_model_key,
-        aug_safety_settings,
-        res_safety_settings,
+        augment_model_name,
+        augment_config,
+        augment_safety_settings,
+        augment_model_api_key,
+        response_model_name,
+        generation_config,
+        response_safety_settings,
+        response_model_api_key,
         chroma_filename,
         chroma_collection_name,
     ):
-        self.aug_model_name = aug_model_name
-        self.res_model_name = res_model_name
-        self.aug_config = aug_config
-        self.res_config = res_config
-        self.aug_model_key = aug_model_key
-        self.res_model_key = res_model_key
-        self.aug_safety_settings = aug_safety_settings
-        self.res_safety_settings = res_safety_settings
-        self.chroma_filename = chroma_filename
-        self.chroma_collection_name = chroma_collection_name
-        self.embeddings = (self._initialize_embeddings_function,)
+        self.augment_model_name = augment_model_name
+        self.augment_config = augment_config
+        self.augment_safety_settings = augment_safety_settings
+        self._configure_generative_ai(response_model_api_key)
+        self._configure_augment_ai(augment_model_api_key)
+        self.response_model = self._initialize_generative_model(
+            response_model_name, generation_config, response_safety_settings
+        )
+        self.embedding_function = self._initialize_embedding_function()
         self.chroma_collection = self._load_chroma(
             chroma_filename, chroma_collection_name
         )
+        self.cross_encoder = self._initialize_cross_encoder()
 
-    def _handle_error(self, message, exception):
-        print(f"{message} : {str(exception)}")
-        logging.error((f"{message} : {str(exception)}"))
+    def _configure_generative_ai(self, response_model_api_key):
+        try:
+            genai.configure(api_key=response_model_api_key)
+        except Exception as e:
+            self._handle_error("Error configuring generative AI module", e)
+
+    def _configure_augment_ai(self, augment_model_api_key):
+        try:
+            palm.configure(api_key=augment_model_api_key)
+        except Exception as e:
+            self._handle_error("Error configuring augmentation AI module", e)
+
+    def _initialize_generative_model(
+        self, response_model_name, generation_config, response_safety_settings
+    ):
+        try:
+            return genai.GenerativeModel(
+                model_name=response_model_name,
+                generation_config=generation_config,
+                safety_settings=response_safety_settings,
+            )
+        except Exception as e:
+            self._handle_error("Error initializing generative model", e)
+
+    def _initialize_augment_model(
+        self, augment_model_name, augment_config, augment_safety_settings
+    ):
+        try:
+            return palm.GenerativeModel(
+                model_name=augment_model_name,
+                generation_config=augment_config,
+                safety_settings=augment_safety_settings,
+            )
+        except Exception as e:
+            self._handle_error("Error initializing augmentation model", e)
+
+    def _initialize_embedding_function(self):
+        try:
+            return GeminiEmbeddingFunction()
+        except Exception as e:
+            self._handle_error("Error initializing embedding function", e)
 
     def _load_chroma(self, chroma_filename, chroma_collection_name):
         try:
@@ -61,16 +103,128 @@ class Brain:
                 embedding_function=self.embedding_function,
             )
         except Exception as e:
-            self._handle_error("Error loading Chroma collection", e)
+            self._handle_error("Error loading chroma collection", e)
+
+    def _initialize_cross_encoder(self):
+        try:
+            return CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
+        except Exception as e:
+            self._handle_error("Error initializing CrossEncoder model", e)
+
+    def _handle_error(self, message, exception):
+        print(f"{message}: {str(exception)}")
+        logging.error(f"{message}: {str(exception)}")
+
+    def generate_alternative_queries(self, query):
+        try:
+            prompt_template = """ Your task is to break down the query in sub questions and turn it into questions in to ten different ways.Keep in mind, Output one query per line, without numbering the queries.\nQUESTION: '{}'\nANSWER:\n"""
+            prompt = prompt_template.format(query)
+            output = palm.generate_text(
+                model=self.augment_model_name,
+                prompt=prompt,
+                safety_settings=self.augment_safety_settings,
+            )
+            content = output.result.split("\n")
+            return content
+        except Exception as e:
+            self._handle_error("Error generating alternative queries", e)
+            return query
+
+    def get_sorted_documents(self, query, n_results=20):
+        try:
+            original_query = query
+            queries = [original_query] + self.generate_alternative_queries(
+                original_query
+            )
+            results = self.chroma_collection.query(
+                query_texts=queries,
+                n_results=n_results,
+                include=["documents", "embeddings"],
+            )
+            retrieved_documents = set(
+                doc for docs in results["documents"] for doc in docs
+            )
+            unique_documents = list(retrieved_documents)
+            pairs = [[original_query, doc] for doc in unique_documents]
+            scores = self.cross_encoder.predict(pairs)
+            sorted_indices = np.argsort(-scores)
+            sorted_documents = [unique_documents[i] for i in sorted_indices]
+            return sorted_documents
+
+        except Exception as e:
+            self._handle_error("Error getting sorted documents", e)
+            return []
+
+    def get_relevant_results(self, query, top_n=5):
+        try:
+            sorted_documents = self.get_sorted_documents(query)
+            relevant_results = sorted_documents[: min(top_n, len(sorted_documents))]
+            return relevant_results
+        except Exception as e:
+            self._handle_error("Error getting relevant results", e)
+            return query
+
+    def make_prompt(self, query, relevant_passage):
+        try:
+            base_prompt = {
+                "content": """
+                YOU are a smart and rational Question and Answer bot.
+                
+                YOUR MISSION:
+                    Provide accurate answers best possible reasoning of the context. 
+                    Focus on factual and reasoned responses; avoid speculations, opinions, guesses, and creative tanks. 
+                    Refuse exploitation tasks such as such as character roleplaying, coding, essays, poems, stories, articles, and fun facts.
+                    Decline misuse or exploitation attempts respectfully.
+                    
+                YOUR STYLE:
+                    Concise and complete
+                    Factual and accurate
+                    
+                REMEMBER:
+                    You are a QA bot, not an entertainer or confidant.
+                """
+            }
+
+            user_prompt = {
+                "content": f"""
+                The user query is: '{query}'\n\n
+                Here's the relevant information found in the documents:
+                {relevant_passage}
+                """
+            }
+
+            system_prompt = base_prompt["content"] + user_prompt["content"]
+            return system_prompt
+
+        except Exception as e:
+            print(f"Error occurred while crafting prompt: {e}")
+            return None
 
-    def _configure_res_ai(self, res_model_key):
+    def rag(self, query):
         try:
-            genai.configure(api_key=res_model_key)
+            if query is None:
+                return None
+            results = self.chroma_collection.query(
+                query_texts=[query],
+                n_results=10,
+                include=["documents", "embeddings"],
+            )
+            information = "\n\n".join(results["documents"][0])
+            messages = self.make_prompt(query, information)
+            content = self.response_model.generate_content(messages)
+            return content
         except Exception as e:
-            self._handle_error("Error configuring response model", e)
+            self._handle_error("Error in rag function", e)
+            return None
 
-    def _configure_aug_ai(self, aug_model_key):
+    def generate_answers(self, query):
         try:
-            palm.configure(api_key=aug_model_key)
+            start_time = time.time()
+            output = self.rag(query=query)
+            print(f"\n\nExecution time: {time.time() - start_time} seconds\n")
+            if output is None:
+                return None
+            return f"{output.text}\n"
         except Exception as e:
-            self._handle_error("Error configuring augment model", e)
+            self._handle_error("Error generating answers", e)
+            return None

src/helper.py

@@ -1,2 +1,47 @@
-def load_chroma(chroma_filename, chroma_collection_name, embedding_function):
-    pass
+from pypdf import PdfReader
+import chromadb
+from langchain.text_splitter import (
+    RecursiveCharacterTextSplitter,
+    SentenceTransformersTokenTextSplitter,
+)
+def _read_pdf(filename):
+    reader = PdfReader(filename)
+    pdf_texts = [p.extract_text().strip() for p in reader.pages]
+    pdf_texts = [text for text in pdf_texts if text]
+    return pdf_texts
+
+
+def _chunk_texts(texts):
+    character_splitter = RecursiveCharacterTextSplitter(
+        separators=["\n\n", "\n", ". ", " ", ""], chunk_size=1600, chunk_overlap=200
+    )
+    character_split_texts = character_splitter.split_text("\n\n".join(texts))
+    token_splitter = SentenceTransformersTokenTextSplitter(
+        chunk_overlap=20, tokens_per_chunk=300
+    )
+    token_split_texts = []
+    for text in character_split_texts:
+        token_split_texts += token_splitter.split_text(text)
+    return token_split_texts
+
+def load_chroma(filename, collection_name, embedding_function):
+    texts = _read_pdf(filename)
+    chunks = _chunk_texts(texts)
+    chroma_client = chromadb.Client()
+    chroma_collection = chroma_client.create_collection(
+        name=collection_name, embedding_function=embedding_function
+    )
+    ids = [str(i) for i in range(len(chunks))]
+    chroma_collection.add(ids=ids, documents=chunks)
+    return chroma_collection
+
+
+def word_wrap(string, n_chars=72):
+    if len(string) < n_chars:
+        return string
+    else:
+        return (
+            string[:n_chars].rsplit(" ", 1)[0]
+            + "\n"
+            + word_wrap(string[len(string[:n_chars].rsplit(" ", 1)[0]) + 1 :], n_chars)
+        )

src/init.py

@@ -0,0 +1,55 @@
+from brain import Brain
+
+
+class Initializer:
+    @staticmethod
+    def initialize(augment_api_key, response_api_key, chroma_filename):
+        response_model_name = "gemini-pro"
+        augment_model_name = "models/text-bison-001"
+        generation_config = {
+            "temperature": 0.9,
+            "top_p": 0.7,
+            "top_k": 1,
+            "max_output_tokens": 2048,
+        }
+        response_safety_settings = [
+            {"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",
+            },
+        ]
+
+        augment_config = {
+            "temperature": 0.9,
+            "top_p": 1,
+            "top_k": 80,
+            "max_output_tokens": 1024,
+        }
+        augment_safety_settings = [
+            {"category": "HARM_CATEGORY_DEROGATORY", "threshold": 4},
+            {"category": "HARM_CATEGORY_TOXICITY", "threshold": 4},
+            {"category": "HARM_CATEGORY_VIOLENCE", "threshold": 4},
+            {"category": "HARM_CATEGORY_SEXUAL", "threshold": 4},
+            {"category": "HARM_CATEGORY_MEDICAL", "threshold": 4},
+            {"category": "HARM_CATEGORY_DANGEROUS", "threshold": 4},
+        ]
+        chroma_collection_name = str.upper(chroma_filename) + "_COLLECT"
+
+        return Brain(
+            augment_model_name,
+            augment_config,
+            augment_safety_settings,
+            augment_api_key,
+            response_model_name,
+            generation_config,
+            response_safety_settings,
+            response_api_key,
+            chroma_filename,
+            chroma_collection_name,
+        )