Update app.py

app.py

@@ -1,9 +1,6 @@
-# Streamlit application
-import streamlit as st
 import os
 from getpass import getpass
-from transformers import pipeline
-
+import streamlit as st
 from llama_index.node_parser import SemanticSplitterNodeParser
 from llama_index.embeddings import OpenAIEmbedding
 from llama_index.ingestion import IngestionPipeline
@@ -14,12 +11,17 @@ from llama_index import VectorStoreIndex
 from llama_index.retrievers import VectorIndexRetriever
 from llama_index.query_engine import RetrieverQueryEngine
 
-# Function to initialize the Pinecone and LlamaIndex setup
-def initialize_pipeline():
-    pinecone_api_key = os.getenv("PINECONE_API_KEY") 
-    openai_api_key = os.getenv("OPENAI_API_KEY") 
+# Streamlit interface for inputting API keys
+st.title("Annual Report Summary Chatbot")
+pinecone_api_key = os.getenv("PINECONE_API_KEY")
+openai_api_key = os.getenv("OPENAI_API_KEY") 
+
 
+if pinecone_api_key and openai_api_key:
+    # This will be the model we use both for Node parsing and for vectorization
     embed_model = OpenAIEmbedding(api_key=openai_api_key)
+
+    # Define the initial pipeline
     pipeline = IngestionPipeline(
         transformations=[
             SemanticSplitterNodeParser(
@@ -31,41 +33,32 @@ def initialize_pipeline():
         ],
     )
 
+    # Initialize connection to Pinecone
     pc = PineconeGRPC(api_key=pinecone_api_key)
     index_name = "anualreport"
+
+    # Initialize your index
     pinecone_index = pc.Index(index_name)
+
+    # Initialize VectorStore
     vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
-    pinecone_index.describe_index_stats()
 
-    if not os.getenv('OPENAI_API_KEY'):
-        os.environ['OPENAI_API_KEY'] = openai_api_key
+    pinecone_index.describe_index_stats()
 
+    # Instantiate VectorStoreIndex object from our vector_store object
     vector_index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
-    retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=5)
-    query_engine = RetrieverQueryEngine(retriever=retriever)
-    
-    return query_engine
-
-# Streamlit UI
-st.title("Chat with Annual Reports")
 
-# Initialize the query engine
-query_engine = initialize_pipeline()
-
-# Conversation model using Hugging Face transformers
-conversation_pipeline = pipeline("conversational", model="microsoft/DialoGPT-medium")
-
-# User input
-user_input = st.text_input("You: ", "")
+    # Grab 5 search results
+    retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=5)
 
-if user_input:
-    # Query the vector DB
-    llm_query = query_engine.query(user_input)
-    response = llm_query.response
+    # Query vector DB
+    query_engine = RetrieverQueryEngine(retriever=retriever)
 
-    # Generate response using Hugging Face conversation model
-    conversation = conversation_pipeline([user_input, response])
-    bot_response = conversation[-1]["generated_text"]
+    user_query = st.text_input("Ask a question about the annual report:")
+    if st.button("Submit"):
+        llm_query = query_engine.query(user_query)
+        st.write(llm_query.response)
 
-    # Display response
-    st.text_area("Bot: ", bot_response, height=200)
+# Run the app
+if __name__ == "__main__":
+    st.run()