Upload 5 files

app.py

@@ -4,19 +4,20 @@ import gradio as gr
 
 from typing import List
 
-from llm_client import PalmClient
+from llm_client import GeminiClient
 from pinecone_index import PinceconeIndex
 
-SYSTEM_MESSAGE = 'Give a precise answer to the question based on only the \
-            context and evidence and do not be verbose.'
-TOP_K = 2
+SYSTEM_MESSAGE = "Give a precise answer to the question based on only the \
+            context and evidence and do not be verbose."
+TOP_K = 1
 
 
 def format_prompt(question: str, evidence: List[str]):
-    evidence_string = ''
+    evidence_string = ""
     for i, ev in enumerate(evidence):
-        evidence_string.join(f'\n Evidence {i+1}: {ev}')
+        evidence_string = evidence_string.join(f"\n Evidence {i+1}: {ev}")
 
+    print(f"evidence string - {evidence_string}")
     content = f"{SYSTEM_MESSAGE} \
               \n ### Question:{question} \
               \n ### Evidence: {evidence_string} \
@@ -25,39 +26,36 @@ def format_prompt(question: str, evidence: List[str]):
     return content
 
 
-if __name__ == '__main__':
-
-    config_path = 'config.yml'
-    with open('config.yml', 'r') as file:
+if __name__ == "__main__":
+    config_path = "config.yml"
+    with open("config.yml", "r") as file:
         config = yaml.safe_load(file)
 
     print(config)
 
-    data_path = config['paths']['data_path']
-    project = config['paths']['project']
+    data_path = config["paths"]["data_path"]
+    project = config["paths"]["project"]
 
-    index_name = config['pinecone']['index-name']
-    embedding_model = config['sentence-transformers']['model-name']
-    embedding_dimension = config['sentence-transformers'][
-        'embedding-dimension']
+    index_name = config["pinecone"]["index-name"]
+    embedding_model = config["sentence-transformers"]["model-name"]
+    embedding_dimension = config["sentence-transformers"]["embedding-dimension"]
 
     index = PinceconeIndex(index_name, embedding_model)
     index.connect_index(embedding_dimension, False)
 
-    palm_client = PalmClient()
+    gemini_client = GeminiClient()
 
     def get_answer(question: str):
         evidence = index.query(question, top_k=TOP_K)
         prompt_with_evidence = format_prompt(question, evidence)
         print(prompt_with_evidence)
-        response = palm_client.generate_text(prompt_with_evidence)
+        response = gemini_client.generate_text(prompt_with_evidence)
         final_output = [response] + evidence
 
         return final_output
 
-    context_outputs = [gr.Textbox(label=f'Evidence {i+1}')
-                       for i in range(TOP_K)]
-    result_output = [gr.Textbox(label='Answer')]
+    context_outputs = [gr.Textbox(label=f"Evidence {i+1}") for i in range(TOP_K)]
+    result_output = [gr.Textbox(label="Answer")]
 
     gradio_outputs = result_output + context_outputs
     gradio_inputs = gr.Textbox(placeholder="Enter your question...")
@@ -69,12 +67,8 @@ if __name__ == '__main__':
         outputs=gradio_outputs,
         title="GT Student Code of Conduct Bot",
         description="Get LLM-powered answers to questions about the \
-                Georgia Tech Student Code of Conduct. The evidence contains exerpts\
-                from the Code of Conduct. LLMs are prone to make mistakes \
-                and this is an experimental setup. Please do not consider the \
-                answers as legal advice. We recommend you to consult the \
-                actual Student Code of Conduct online for authoritative \
-                information."
+            Georgia Tech Student Code of Conduct. The evidences are exerpts\
+                from the Code of Conduct.",
     )
 
     demo.launch()

llm_client.py

@@ -1,41 +1,45 @@
 import os
 
-import google.generativeai as palm
+import google.generativeai as genai
 
 
-class PalmClient:
+class GeminiClient:
     def __init__(self):
         self.connect_client()
 
     def connect_client(self):
-        if (not os.getenv('GOOGLE_PALM_KEY')):
-            raise Exception('Please set your Google MakerSuite API key')
-
-        api_key = os.getenv('GOOGLE_PALM_KEY')
-        palm.configure(api_key=api_key)
-
-        safety_overrides = [
-            {"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}
-            ]
+        if not os.getenv("GOOGLE_PALM_KEY"):
+            raise Exception("Please set your Google AI Studio key")
+
+        api_key = os.getenv("GOOGLE_PALM_KEY")
+        genai.configure(api_key=api_key)
+
+        safety_settings = [
+            {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"},
+            {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_ONLY_HIGH"},
+            {
+                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                "threshold": "BLOCK_ONLY_HIGH",
+            },
+            {
+                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                "threshold": "BLOCK_ONLY_HIGH",
+            },
+        ]
 
         defaults = {
-            'model': 'models/text-bison-001',
-            'temperature': 0.7,
-            'candidate_count': 1,
-            'top_k': 40,
-            'top_p': 0.95,
-            'max_output_tokens': 1024,
-            'stop_sequences': [],
-            'safety_settings': safety_overrides,
+            "temperature": 0.7,
+            "top_k": 40,
+            "top_p": 0.95,
+            "max_output_tokens": 1024,
         }
 
-        self.defaults = defaults
+        self.model = genai.GenerativeModel(
+            model_name="gemini-pro",
+            generation_config=defaults,
+            safety_settings=safety_settings,
+        )
 
     def generate_text(self, prompt: str) -> str:
-        response = palm.generate_text(**self.defaults, prompt=prompt)
-        return response.candidates[0]['output']
+        response = self.model.generate_content(prompt)
+        return response.text

pinecone_index.py

@@ -19,18 +19,17 @@ class PinceconeIndex:
         self.index_name = index_name
         self._embeddingModel = HuggingFaceEmbeddings(model_name=model_name)
 
-    def connect_index(self, embedding_dimension: int,
-                      delete_existing: bool = False):
+    def connect_index(self, embedding_dimension: int, delete_existing: bool = False):
         index_name = self.index_name
 
         # load pinecone env variables within Google Colab
-        if (not os.getenv('PINECONE_KEY')) or (not os.getenv('PINECONE_ENV')):
-            dotenv_path = Path('/content/gt-policy-bot/config.env')
+        if (not os.getenv("PINECONE_KEY")) or (not os.getenv("PINECONE_ENV")):
+            dotenv_path = Path("/content/gt-policy-bot/config.env")
             load_dotenv(dotenv_path=dotenv_path)
 
         pinecone.init(
-            api_key=os.getenv('PINECONE_KEY'),
-            environment=os.getenv('PINECONE_ENV'),
+            api_key=os.getenv("PINECONE_KEY"),
+            environment=os.getenv("PINECONE_ENV"),
         )
 
         if index_name in pinecone.list_indexes() and delete_existing:
@@ -47,9 +46,8 @@ class PinceconeIndex:
     def upsert_docs(self, df: pd.DataFrame, text_col: str):
         loader = DataFrameLoader(df, page_content_column=text_col)
         docs = loader.load()
-        Pinecone.from_documents(docs, self._embeddingModel,
-                                index_name=self.index_name)
-    
+        Pinecone.from_documents(docs, self._embeddingModel, index_name=self.index_name)
+
     def get_embedding_model(self):
         return self._embeddingModel
 
@@ -57,47 +55,44 @@ class PinceconeIndex:
         return self.index_name
 
     def query(self, query: str, top_k: int = 5) -> List[str]:
-        docsearch = Pinecone.from_existing_index(self.index_name,
-                                                 self._embeddingModel)
+        docsearch = Pinecone.from_existing_index(self.index_name, self._embeddingModel)
         res = docsearch.similarity_search(query, k=top_k)
 
         return [doc.page_content for doc in res]
 
 
-if __name__ == '__main__':
-    config_path = 'config.yml'
-    with open('config.yml', 'r') as file:
+if __name__ == "__main__":
+    config_path = "config.yml"
+    with open("config.yml", "r") as file:
         config = yaml.safe_load(file)
 
     print(config)
 
-    data_path = config['paths']['data_path']
-    project = config['paths']['project']
-    format = '.csv'
+    data_path = config["paths"]["data_path"]
+    project = config["paths"]["project"]
+    format = ".csv"
 
-    index_name = config['pinecone']['index-name']
-    embedding_model = config['sentence-transformers'][
-        'model-name']
-    embedding_dimension = config['sentence-transformers'][
-        'embedding-dimension']
+    index_name = config["pinecone"]["index-name"]
+    embedding_model = config["sentence-transformers"]["model-name"]
+    embedding_dimension = config["sentence-transformers"]["embedding-dimension"]
     delete_existing = True
 
-    if config['paths']['chunking'] == 'manual':
+    if config["paths"]["chunking"] == "manual":
         print("Using manual chunking")
-        file_path_embedding = config['paths']['manual_chunk_file']
-        df = pd.read_csv(file_path_embedding, header=None, names=['chunks'])
+        file_path_embedding = config["paths"]["manual_chunk_file"]
+        df = pd.read_csv(file_path_embedding, header=None, names=["chunks"])
     else:
         print("Using automatic chunking")
-        file_path_embedding = config['paths']['auto_chunk_file']
+        file_path_embedding = config["paths"]["auto_chunk_file"]
         df = pd.read_csv(file_path_embedding, index_col=0)
 
     print(df)
     start_time = time.time()
     index = PinceconeIndex(index_name, embedding_model)
     index.connect_index(embedding_dimension, delete_existing)
-    index.upsert_docs(df, 'chunks')
+    index.upsert_docs(df, "chunks")
     end_time = time.time()
-    print(f'Indexing took {end_time - start_time} seconds')
+    print(f"Indexing took {end_time - start_time} seconds")
 
     index = PinceconeIndex(index_name, embedding_model)
     index.connect_index(embedding_dimension, delete_existing=False)