Update app.py

app.py

@@ -44,26 +44,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)
-#     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 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)
-    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
+# # 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)
+#     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
@@ -75,12 +75,12 @@ 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")
+# # 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:
@@ -99,11 +99,26 @@ if prompt := st.chat_input("Enter your query"):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        response = chat_with_google_llm(prompt, vectordb, tokenizer, model)
+        openai_api_key = st.secrets["OPENAI_API_KEY"]
+        response = chat_with_openai(prompt, vectordb, openai_api_key)
         st.markdown(response)
         
     st.session_state.messages.append({"role": "assistant", "content": response})
+
+# User input
+# if prompt := st.chat_input("Enter your query"):
+#     st.session_state.messages.append({"role": "user", "content": prompt})
+#     with st.chat_message("user"):
+#         st.markdown(prompt)
+
+#     with st.chat_message("assistant"):
+#         response = chat_with_google_llm(prompt, vectordb, tokenizer, model)
+#         st.markdown(response)
         
+#     st.session_state.messages.append({"role": "assistant", "content": response})
+
+
+    
 # # Query input
 # query = st.text_input("Enter your query", "")