Update app.py

app.py

@@ -43,17 +43,26 @@ def augment_prompt(query, vectordb):
     return augmented_prompt
 
 
-# Function to handle chat with OpenAI
-def chat_with_openai(query, vectordb, openai_api_key):
-    chat = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=openai_api_key)
+# # Function to handle chat with OpenAI
+# def chat_with_openai(query, vectordb, openai_api_key):
+#     chat = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=openai_api_key)
+#     augmented_query = augment_prompt(query, vectordb)
+#     prompt = HumanMessage(content=augmented_query)
+#     messages = [
+#         SystemMessage(content="You are a helpful assistant."),
+#         prompt
+#     ]
+#     res = chat(messages)
+#     return res.content
+
+
+# Function to handle chat with the Google open-source LLM
+def chat_with_google_llm(query, vectordb, tokenizer, model):
     augmented_query = augment_prompt(query, vectordb)
-    prompt = HumanMessage(content=augmented_query)
-    messages = [
-        SystemMessage(content="You are a helpful assistant."),
-        prompt
-    ]
-    res = chat(messages)
-    return res.content
+    input_ids = tokenizer(augmented_query, return_tensors="pt") #.to("cuda")
+    outputs = model.generate(input_ids, max_length=512, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
 
 # Streamlit UI
@@ -65,6 +74,13 @@ zip_file_path = "chroma_db_compressed_.zip"
 extract_path = "./chroma_db_extracted"
 vectordb = load_vector_db(zip_file_path, extract_path)
 
+# Load Google model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
+model = AutoModelForCausalLM.from_pretrained(
+    "google/gemma-2b-it",
+    torch_dtype=torch.bfloat16
+)#.to("cuda")
+
 # Initialize session state for chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
@@ -82,8 +98,7 @@ if prompt := st.chat_input("Enter your query"):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        openai_api_key = st.secrets["OPENAI_API_KEY"]
-        response = chat_with_openai(prompt, vectordb, openai_api_key)
+        response = chat_with_google_llm(prompt, vectordb, tokenizer, model)
         st.markdown(response)
         
     st.session_state.messages.append({"role": "assistant", "content": response})