Update app.py

app.py

@@ -1,6 +1,5 @@
-#import atexit
+import atexit
 import gradio as gr
-from gradio import Button
 #from langchain.document_loaders import UnstructuredPDFLoader
 from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -86,6 +85,11 @@ def generate_random_string(length):
     return ''.join(random.choice(letters) for i in range(length))      
 random_string = generate_random_string(10)
 
+def exit_handler():
+    pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
+    index_namespace_to_delete = pinecone.Index(index_name=index_name)
+    index_namespace_to_delete.delete(delete_all=True, namespace=namespace)   
+
 pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
 index_name = PINECONE_INDEX_NAME
 #index_name = pinecone.Index(index_name)
@@ -128,33 +132,35 @@ PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "q
 
 chain = load_qa_chain(llm=llm, chain_type="stuff", prompt=PROMPT)
 
-#chain = load_qa_chain(llm=llm, chain_type="stuff")  
+#chain = load_qa_chain(llm=llm, chain_type="stuff")
 
-def run_chain(delete_button, user_query):
+def run_chain(user_query):
     pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
     if user_query !="" and not user_query.strip().isspace() and not user_query.isspace():
-        print("Your query:\n"+user_query)
-        vector_db_from_index = Pinecone.from_existing_index(index_name, hf_embeddings, namespace=namespace)
-        ss_results = vector_db_from_index.similarity_search(query=user_query, namespace=namespace, k=5)
-        initial_ai_response = chain.run(input_documents=ss_results, question=user_query, return_only_outputs=True)
-        temp_ai_response = initial_ai_response.partition('<|end|>')[0]
-        final_ai_response = temp_ai_response.replace('\n', '')
-        print(final_ai_response)
-        print(index_status)
-        print(index_name_extracted)
-        print(namespace)
-        print("****************")
-        return final_ai_response
+      print("Your query:\n"+user_query)
+      vector_db_from_index = Pinecone.from_existing_index(index_name, hf_embeddings, namespace=namespace)
+      ss_results = vector_db_from_index.similarity_search(query=user_query, namespace=namespace, k=5)
+      initial_ai_response = chain.run(input_documents=ss_results, question=user_query, return_only_outputs=True)        
+      #initial_ai_response=chain({"input_documents": ss_results, "question": user_query}, return_only_outputs=True)            
+      temp_ai_response = initial_ai_response.partition('<|end|>')[0]
+      final_ai_response = temp_ai_response.replace('\n', '')
+      print(final_ai_response)
+      print(index_status)
+      print(index_name_extracted)
+      print(namespace) 
+      print("****************")
+      return final_ai_response
     else:
-        print("Invalid inputs.")
+      print("Invalid inputs.")  
 
-def delete_index_namespace():
-    pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
-    index_namespace_to_delete = pinecone.Index(index_name=index_name)
-    index_namespace_to_delete.delete(delete_all=True, namespace=namespace)
+#def delete_index_namespace():
+#    pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
+#    index_namespace_to_delete = pinecone.Index(index_name=index_name)
+#    index_namespace_to_delete.delete(delete_all=True, namespace=namespace)
+
+iface = gr.Interface(fn=run_chain, inputs="text", outputs="text", title="AI Response")
+iface.launch()
 
-delete_button = gr.Button("Delete Index Namespace", callback=delete_index_namespace)
-#delete_button = gr.Button("Delete Index Namespace", delete_index_namespace)
-#delete_button = gr.Button("Delete Index Namespace", description="点击此按钮删除索引命名空间")
-iface = gr.Interface(fn=run_chain, inputs=[delete_button, "text"], outputs="text", title="AI Response")
-iface.launch()
+delete_button = Button("Exit & Clear Data")
+if delete_button:
+    atexit.register(exit_handler)