gemini

app.py

@@ -1,23 +1,25 @@
 # AI assistant with a RAG system to query information from
-#  the gwIAS search pipeline
+#  the gwIAS search pipline
 # using Langchain and deployed with Gradio
 
 from rag import RAG, load_docs
+from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_community.embeddings import HuggingFaceInstructEmbeddings
-from langchain_community.chat_models import ChatOpenAI
+#from langchain.chat_models import ChatOpenAI
 import gradio as gr
+import os
 
 # Load the documentation
 docs = load_docs()
 print("Pages loaded:", len(docs))
 
 # LLM model
-llm = ChatOpenAI(model="gpt-4") # Fixed model name to use a real model
+#llm = ChatOpenAI(model="gpt-4o-mini")
+llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
 
 # Embeddings
 embed_model = "sentence-transformers/multi-qa-distilbert-cos-v1"
 # embed_model = "nvidia/NV-Embed-v2"
-# text-embedding-3-small
 embeddings = HuggingFaceInstructEmbeddings(model_name=embed_model)
 
 # RAG chain
@@ -36,9 +38,10 @@ def handle_prompt(message, history):
 
 
 if __name__ == "__main__":
+
     # Predefined messages and examples
     description = "AI powered assistant to help with [gwIAS](https://github.com/JayWadekar/gwIAS-HM) gravitational wave search pipeline."
-    greeting_message = "Hi, I'm the gwIAS Bot, I'm here to assist you with the search pipeline."
+    greetingsmessage = "Hi, I'm the gwIAS Bot, I'm here to assist you with the search pipeline."
     example_questions = [
         "Can you give me the code for calculating coherent score?",
         "Which module in the code is used for collecting coincident triggers?",
@@ -46,23 +49,24 @@ if __name__ == "__main__":
     ]
 
     # Define customized Gradio chatbot
-    chatbot = gr.Chatbot(
-        [{"role": "assistant", "content": greeting_message}],
-        type="messages",
-        avatar_images=["ims/userpic.png", "ims/gwIASlogo.jpg"],
-        height="60vh"
-    )
+    chatbot = gr.Chatbot([{"role": "assistant", "content": greetingsmessage}],
+                         type="messages",
+                         avatar_images=["ims/userpic.png", "ims/gwIASlogo.jpg"],
+                         height="60vh")
 
     # Define Gradio interface
-    demo = gr.ChatInterface(
-        fn=handle_prompt,
-        chatbot=chatbot,
-        title="gwIAS DocBot",
-        description=description,
-        examples=example_questions,
-        theme=gr.themes.Soft(),
-        fill_height=True
-    )
+    demo = gr.ChatInterface(handle_prompt,
+                            type="messages",
+                            title="gwIAS DocBot",
+                            fill_height=True,
+                            examples=example_questions,
+                            theme=gr.themes.Soft(),
+                            description=description,
+                            # cache_examples=False,
+                            chatbot=chatbot)
 
-    # Launch the interface
     demo.launch()
+
+# https://arxiv.org/html/2405.17400v2
+# https://arxiv.org/html/2312.06631v1
+# https://arxiv.org/html/2310.15233v2