Update pages/bot.py

pages/bot.py

@@ -7,6 +7,11 @@ import os
 from PyPDF2 import PdfReader
 from transformers import pipeline
 
+#Retriever erweiterung
+from langchain.prompts import ChatPromptTemplate
+from langchain.schema import StrOutputParser
+from langchain.schema.runnable import RunnablePassthrough
+
 ###########
 #pip install faiss-cpu
 #pip install langchain
@@ -82,7 +87,9 @@ def main():
     if user_question:
         st.text(retrieved_docs[0].page_content)
         context=retrieved_docs[0].page_content
-        #question=user_question
+        question=user_question
+
+        ##IDEE Text Generation
         generator = pipeline('text-generation', model = 'gpt2')
         answer = generator(context, max_length = 30, num_return_sequences=3)
         
@@ -90,12 +97,30 @@ def main():
         #st.text_area()
         st.text(answer)
         st.text(type(answer))
-                  # bei incoming pdf
 
-    #vectorstore_DB=get_vectorstore()        # bei Abfrage durch Chatbot
-    #print(get_vectorstore().similarity_search_with_score("stelle")) # zeigt an ob Vektordatenbank gefüllt ist
-    
-    #print(get_conversation_chain(get_vectorstore()))
+        #IDEE Retriever erweitern
+        template = """Answer the question based only on the following context:
+
+        {context}
+        
+        Question: {question}
+        """
+        prompt = ChatPromptTemplate.from_template(template)
+        model = "hkunlp/instructor-base"
+        
+        
+        def format_docs(docs):
+            return "\n\n".join([d.page_content for d in docs])
+        
+        
+        chain = (
+            {"context": retriever | format_docs, "question": RunnablePassthrough()}
+            | prompt
+            | model
+            | StrOutputParser()
+        )
+
+        st.text(chain.invoke(question))