Update app.py

app.py

@@ -15,50 +15,56 @@ st.set_page_config(page_title="Chat with Notes and AI", page_icon=":books:", lay
 # Load environment variables
 load_dotenv()
 
-# Optimized pipeline setup
+# Dolly-v2-3b model pipeline
 @st.cache_resource
 def load_pipeline():
-    # Use a smaller model for faster performance
-    model_name = "databricks/dolly-v2-1b"  # Switch to a lighter model
-    
-    # Load tokenizer and model
+    model_name = "databricks/dolly-v2-3b"
+
+    # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left", trust_remote_code=True)
+
+    # Load model with offload folder for disk storage of weights
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,  # Use float32 for CPU compatibility
-        device_map="auto",         # Automatically map devices
-        trust_remote_code=True
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,  # Use float16 for GPU, float32 for CPU
+        device_map="auto",          # Automatically map model to available devices (e.g., GPU if available)
+        trust_remote_code=True,
+        offload_folder="./offload_weights"  # Folder to store offloaded weights
     )
 
-    # Return text-generation pipeline with full-text output
+    # Return text-generation pipeline
     return pipeline(
         task="text-generation",
         model=model,
         tokenizer=tokenizer,
-        torch_dtype=torch.float32,  # Ensure compatibility with CPU
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         device_map="auto",
-        return_full_text=True,
-        max_new_tokens=100  # Limit response length
+        return_full_text=True
     )
 
-# Initialize pipeline
+# Initialize Dolly pipeline
 generate_text = load_pipeline()
 
-# LangChain Integration
+# Create a HuggingFace pipeline wrapper for LangChain
 hf_pipeline = HuggingFacePipeline(pipeline=generate_text)
 
-# Templates for prompts
-prompt = PromptTemplate(input_variables=["instruction"], template="{instruction}")
+# Template for instruction-only prompts
+prompt = PromptTemplate(
+    input_variables=["instruction"],
+    template="{instruction}"
+)
+
+# Template for prompts with context
 prompt_with_context = PromptTemplate(
     input_variables=["instruction", "context"],
     template="{instruction}\n\nInput:\n{context}"
 )
 
-# LangChain LLM chains
+# Create LLM chains
 llm_chain = LLMChain(llm=hf_pipeline, prompt=prompt)
 llm_context_chain = LLMChain(llm=hf_pipeline, prompt=prompt_with_context)
 
-# Extract content from .txt files
+# Extracting text from .txt files
 def get_text_files_content(folder):
     text = ""
     for filename in os.listdir(folder):
@@ -67,92 +73,107 @@ def get_text_files_content(folder):
                 text += file.read() + "\n"
     return text
 
-# Convert text into chunks for vectorization
+# Converting text to chunks
 def get_chunks(raw_text):
     from langchain.text_splitter import CharacterTextSplitter
     text_splitter = CharacterTextSplitter(
         separator="\n",
-        chunk_size=500,  # Smaller chunks for faster processing
-        chunk_overlap=50  # Minimal overlap
+        chunk_size=1000,  # Reduced chunk size for faster processing
+        chunk_overlap=200,  # Smaller overlap for efficiency
+        length_function=len
     )
-    return text_splitter.split_text(raw_text)
+    chunks = text_splitter.split_text(raw_text)
+    return chunks
 
-# Create FAISS vectorstore for embeddings
+# Using Hugging Face embeddings model and FAISS to create vectorstore
 def get_vectorstore(chunks):
     embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2",  # Lightweight embeddings
-        model_kwargs={'device': 'cpu'}  # Ensure embeddings run on CPU
+        model_name="sentence-transformers/all-MiniLM-L6-v2",
+        model_kwargs={'device': 'cpu'}  # Ensure embeddings use CPU
     )
-    return FAISS.from_texts(texts=chunks, embedding=embeddings)
+    vectorstore = FAISS.from_texts(texts=chunks, embedding=embeddings)
+    return vectorstore
 
-# Handle user queries
+# Generating response from user queries
 def handle_question(question, vectorstore=None):
     if vectorstore:
-        # Retrieve the most relevant chunk
-        documents = vectorstore.similarity_search(question, k=1)  # Retrieve fewer chunks
-        context = "\n".join([doc.page_content for doc in documents])[:500]  # Short context for efficiency
+        # Reduce the number of retrieved chunks for faster processing
+        documents = vectorstore.similarity_search(question, k=2)
+        context = "\n".join([doc.page_content for doc in documents])
+
+        # Limit context to 1000 characters to speed up model inference
+        context = context[:1000]
 
         if context:
-            return llm_context_chain.predict(instruction=question, context=context).strip()
+            result_with_context = llm_context_chain.invoke({"instruction": question, "context": context})
+            return result_with_context
 
-    # Fallback to instruction-only chain if no context
-    return llm_chain.predict(instruction=question).strip()
+    # Fallback to instruction-only chain if no context is found
+    return llm_chain.invoke({"instruction": question})
 
 def main():
     st.title("Chat with Notes :books:")
 
-    # Session state for vectorstore
+    # Initialize session state
     if "vectorstore" not in st.session_state:
         st.session_state.vectorstore = None
 
-    # Data folders
-    data_folder = "data"  # Folder for Current Affairs
-    essay_folder = "essays"  # Folder for Essays
+    # Define folders for Current Affairs and Essays
+    data_folder = "data"  # Current Affairs folders
+    essay_folder = "essays"  # Essays folder
 
     # Content type selection
     content_type = st.sidebar.radio("Select Content Type:", ["Current Affairs", "Essays"])
 
-    # Subjects based on content type
+    # Handle Current Affairs (each subject has its own folder)
     if content_type == "Current Affairs":
-        subjects = [f for f in os.listdir(data_folder) if os.path.isdir(os.path.join(data_folder, f))] if os.path.exists(data_folder) else []
-    else:
-        subjects = [f.replace(".txt", "") for f in os.listdir(essay_folder) if f.endswith('.txt')] if os.path.exists(essay_folder) else []
+        if os.path.exists(data_folder):
+            subjects = [f for f in os.listdir(data_folder) if os.path.isdir(os.path.join(data_folder, f))]
+        else:
+            subjects = []
+    # Handle Essays (all essays are in a single folder)
+    elif content_type == "Essays":
+        if os.path.exists(essay_folder):
+            subjects = [f.replace(".txt", "") for f in os.listdir(essay_folder) if f.endswith('.txt')]
+        else:
+            subjects = []
 
     # Subject selection
     selected_subject = st.sidebar.selectbox("Select a Subject:", subjects)
 
-    # Load content based on selection
+    # Process selected subject
     raw_text = ""
     if content_type == "Current Affairs" and selected_subject:
         subject_folder = os.path.join(data_folder, selected_subject)
         raw_text = get_text_files_content(subject_folder)
     elif content_type == "Essays" and selected_subject:
-        subject_file = os.path.join(essay_folder, f"{selected_subject}.txt")
+        subject_file = os.path.join(essay_folder, selected_subject + ".txt")
         if os.path.exists(subject_file):
             with open(subject_file, "r", encoding="utf-8") as file:
                 raw_text = file.read()
 
-    # Display preview and create vectorstore
+    # Display preview of notes
     if raw_text:
         st.subheader("Preview of Notes")
-        st.text_area("Preview Content:", value=raw_text[:1000], height=300, disabled=True)
+        st.text_area("Preview Content:", value=raw_text[:2000], height=300, disabled=True)  # Show a snippet of the notes
 
-        if "vectorstore" not in st.session_state or st.session_state.vectorstore is None:
-            chunks = get_chunks(raw_text)
-            st.session_state.vectorstore = get_vectorstore(chunks)
+        # Create vectorstore for Current Affairs or Essays
+        text_chunks = get_chunks(raw_text)
+        vectorstore = get_vectorstore(text_chunks)
+        st.session_state.vectorstore = vectorstore
     else:
         st.warning("No content available for the selected subject.")
 
-    # Question and response
+    # Chat interface
     st.subheader("Ask Your Question")
     question = st.text_input("Ask a question about your selected subject:")
     if question:
         if st.session_state.vectorstore:
             response = handle_question(question, st.session_state.vectorstore)
             st.subheader("Answer:")
-            st.write(response or "No response found.")
+            st.write(response.get("text", "No response found."))
         else:
             st.warning("Please load the content for the selected subject before asking a question.")
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()