Update app.py

app.py

@@ -5,7 +5,7 @@ import torch
 import logging
 
 from operator import itemgetter
-from langchain_openai import ChatOpenAI
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_core.prompts import ChatPromptTemplate
@@ -13,9 +13,6 @@ from langchain_community.vectorstores.chroma import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.schema import AIMessage, HumanMessage
 from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnableLambda, RunnablePassthrough
-from langchain.chains.combine_documents import create_stuff_documents_chain
-from langchain.chains import create_retrieval_chain
 from langchain.globals import set_debug
 from dotenv import load_dotenv
 
@@ -26,16 +23,27 @@ set_debug(True)
 load_dotenv()
 
 openai_api_key = os.getenv("OPENAI_API_KEY")
+langchain_api_key = os.getenv("LANGCHAIN_API_KEY")
+langchain_endpoint = os.getenv("LANGCHAIN_ENDPOINT")
+langchain_project_id = os.getenv("LANGCHAIN_PROJECT")
+access_key = os.getenv("ACCESS_TOKEN_SECRET")
 
 persist_dir = "./chroma_db"
-device='cuda:0'
-model_name="all-mpnet-base-v2"
-model_kwargs = {'device': device if torch.cuda.is_available() else 'cpu'}
+device = 'cuda:0'
+model_name = "all-mpnet-base-v2"
+model_kwargs = {'device': device if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"}
 logging.info(f"Using device {model_kwargs['device']}")
-# Create embeddings and store in vectordb
-embeddings = HuggingFaceEmbeddings(model_name=model_name, show_progress=True, model_kwargs=model_kwargs)
+embed_money = False
 
-def configure_retriever(local_files, chunk_size=12500, chunk_overlap=2500):
+# Create embeddings and store in vectordb
+if embed_money:
+    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
+    logging.info(f"Using OpenAI embeddings")
+else:
+    embeddings = HuggingFaceEmbeddings(model_name=model_name, show_progress=True, model_kwargs=model_kwargs)
+    logging.info(f"Using HuggingFace embeddings")
+
+def configure_retriever(local_files, chunk_size=15000, chunk_overlap=2500):
     logging.info("Configuring retriever")
     
     if not os.path.exists(persist_dir):
@@ -63,10 +71,8 @@ def configure_retriever(local_files, chunk_size=12500, chunk_overlap=2500):
         vectordb = Chroma.from_documents(splits, embeddings, persist_directory=persist_dir)
 
         # Define retriever
-        retriever = vectordb.as_retriever(
-                search_type="similarity_score_threshold",
-                search_kwargs={'score_threshold': 0.8}
-            )
+        retriever = vectordb.as_retriever(search_type="mmr", search_kwargs={'k': 6, 'lambda_mult': 0.25})
+
 
         return retriever
     else:
@@ -74,10 +80,7 @@ def configure_retriever(local_files, chunk_size=12500, chunk_overlap=2500):
         vectordb = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
 
         # Define retriever
-        retriever = vectordb.as_retriever(
-                search_type="similarity_score_threshold",
-                search_kwargs={'score_threshold': 0.8}
-            )
+        retriever = vectordb.as_retriever(search_type="mmr", search_kwargs={'k': 6, 'lambda_mult': 0.25})
 
         return retriever
 
@@ -86,7 +89,11 @@ local_files = [f for f in os.listdir(directory) if f.endswith(".pdf")]
 
 # Setup LLM
 llm = ChatOpenAI(
-    model_name="gpt-3.5-turbo", openai_api_key=openai_api_key, temperature=0, streaming=True
+    model_name="gpt-4-0125-preview", openai_api_key=openai_api_key, temperature=0.1, streaming=True
+)
+
+llm_translate = ChatOpenAI(
+    model_name="gpt-3.5-turbo", openai_api_key=openai_api_key, temperature=0.0
 )
 
 retriever = configure_retriever(local_files)
@@ -96,7 +103,7 @@ template = """Answer the question based only on the following context:
 
 Question: {question}
 
-Answer in German language.
+Answer in German Language. If the question is not related to the context, answer with "I don't know" in German.
 """
 
 prompt = ChatPromptTemplate.from_template(template)
@@ -111,28 +118,44 @@ chain = (
     | StrOutputParser()
 )
 
+chain_translate = (llm_translate
+    | StrOutputParser()
+)
+
+
 def predict(message, history):
-    message = f"Translate the following text to German: {message}"
+    message = chain_translate.invoke(f"Translate this sentence to English: {message}")
     history_langchain_format = []
     for human, ai in history:
         history_langchain_format.append(HumanMessage(content=human))
         history_langchain_format.append(AIMessage(content=ai))
     history_langchain_format.append(HumanMessage(content=message))
     gpt_response = llm(history_langchain_format)
-    return chain.invoke({"question": gpt_response.content})
+    for chunk in chain.stream({"question": gpt_response.content}): # Stream the response
+        yield chunk
+
+
+image_path = "./ui/logo.png" if os.path.exists("./ui/logo.png") else "./logo.png"
 
-demo = gr.ChatInterface(
+with gr.Blocks() as demo:
+    gr.Image(image_path)
+    gr.ChatInterface(
     predict,
     chatbot=gr.Chatbot(height=500, show_share_button=True),
     textbox=gr.Textbox(placeholder="stell mir Fragen", container=False, scale=7),
     title="Beitrag Service",
     description="Ich bin Ihr hilfreicher KI-Assistent",
     theme="soft",
-    examples=["Hello"],
+    examples=[
+    "Generate auditing questions about Change Management",
+    "Generate auditing questions about Software Maintenance",
+    "Generate auditing questions about Data Protection",
+    "Generate auditing questions about IT",
+    "Generate auditing questions about control systems",
+    "Generate auditing questions about GDPR compliance",
+    ],
     cache_examples=True,
-    retry_btn="Wiederholen",
-    undo_btn="Vorheriges löschen",
-    clear_btn="Löschen").launch(show_api= False)
+    ).launch(show_api= False)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()