Update app.py

app.py

@@ -15,46 +15,50 @@ st.set_page_config(page_title="Chat with Notes and AI", page_icon=":books:", lay
 # Load environment variables
 load_dotenv()
 
-# Optimized Dolly-v2 model pipeline
+# Optimized pipeline setup
 @st.cache_resource
 def load_pipeline():
-    model_name = "databricks/dolly-v2-1b"  # Smaller model for CPU
-
+    # Use a smaller model for faster performance
+    model_name = "databricks/dolly-v2-1b"  # Switch to a lighter model
+    
     # Load tokenizer and model
     tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left", trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,  # Use float32 for CPU
-        device_map="auto",
-        trust_remote_code=True,
-        offload_folder="./offload_weights"  # Folder to store weights if needed
+        torch_dtype=torch.float32,  # Use float32 for CPU compatibility
+        device_map="auto",         # Automatically map devices
+        trust_remote_code=True
     )
 
-    # Create text-generation pipeline
+    # Return text-generation pipeline with full-text output
     return pipeline(
         task="text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_new_tokens=50,  # Limit response length for speed
-        return_full_text=False,
-        device_map="auto"
+        torch_dtype=torch.float32,  # Ensure compatibility with CPU
+        device_map="auto",
+        return_full_text=True,
+        max_new_tokens=100  # Limit response length
     )
 
-# Initialize Dolly pipeline
+# Initialize pipeline
 generate_text = load_pipeline()
 
-# Create HuggingFace pipeline wrapper for LangChain
+# LangChain Integration
 hf_pipeline = HuggingFacePipeline(pipeline=generate_text)
 
-# Prompt templates
+# Templates for prompts
 prompt = PromptTemplate(input_variables=["instruction"], template="{instruction}")
-prompt_with_context = PromptTemplate(input_variables=["instruction", "context"], template="{instruction}\n\nInput:\n{context}")
+prompt_with_context = PromptTemplate(
+    input_variables=["instruction", "context"],
+    template="{instruction}\n\nInput:\n{context}"
+)
 
-# Create LLM chains
+# LangChain LLM chains
 llm_chain = LLMChain(llm=hf_pipeline, prompt=prompt)
 llm_context_chain = LLMChain(llm=hf_pipeline, prompt=prompt_with_context)
 
-# Extract text from .txt files
+# Extract content from .txt files
 def get_text_files_content(folder):
     text = ""
     for filename in os.listdir(folder):
@@ -63,93 +67,92 @@ def get_text_files_content(folder):
                 text += file.read() + "\n"
     return text
 
-# Convert text into smaller chunks
+# Convert text into chunks for vectorization
 def get_chunks(raw_text):
     from langchain.text_splitter import CharacterTextSplitter
     text_splitter = CharacterTextSplitter(
         separator="\n",
-        chunk_size=512,  # Smaller chunks for faster processing
-        chunk_overlap=50,  # Minimal overlap
-        length_function=len
+        chunk_size=500,  # Smaller chunks for faster processing
+        chunk_overlap=50  # Minimal overlap
     )
     return text_splitter.split_text(raw_text)
 
-# Create FAISS vectorstore
+# Create FAISS vectorstore for embeddings
 def get_vectorstore(chunks):
     embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/all-MiniLM-L6-v2",
-        model_kwargs={'device': 'cpu'}  # Force CPU usage for embeddings
+        model_name="sentence-transformers/all-MiniLM-L6-v2",  # Lightweight embeddings
+        model_kwargs={'device': 'cpu'}  # Ensure embeddings run on CPU
     )
     return FAISS.from_texts(texts=chunks, embedding=embeddings)
 
-# Generate response from user queries
+# Handle 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])[:512]  # Shorter context
+        context = "\n".join([doc.page_content for doc in documents])[:500]  # Short context for efficiency
 
         if context:
-            result_with_context = llm_context_chain.invoke({"instruction": question, "context": context})
-            return result_with_context
+            return llm_context_chain.predict(instruction=question, context=context).strip()
 
-    # Fallback to instruction-only chain if no context is found
-    return llm_chain.invoke({"instruction": question})
+    # Fallback to instruction-only chain if no context
+    return llm_chain.predict(instruction=question).strip()
 
 def main():
     st.title("Chat with Notes :books:")
 
-    # Initialize session state
+    # Session state for vectorstore
     if "vectorstore" not in st.session_state:
         st.session_state.vectorstore = None
 
-    # Define folders for Current Affairs and Essays
-    data_folder = "data"  # Current Affairs folders
-    essay_folder = "essays"  # Essays folder
+    # Data folders
+    data_folder = "data"  # Folder for Current Affairs
+    essay_folder = "essays"  # Folder for Essays
 
-    # Sidebar for content selection
+    # Content type selection
     content_type = st.sidebar.radio("Select Content Type:", ["Current Affairs", "Essays"])
 
-    # Handle folder-based selection
+    # Subjects based on content type
     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 []
-    elif content_type == "Essays":
+    else:
         subjects = [f.replace(".txt", "") for f in os.listdir(essay_folder) if f.endswith('.txt')] if os.path.exists(essay_folder) else []
 
+    # Subject selection
     selected_subject = st.sidebar.selectbox("Select a Subject:", subjects)
 
-    # Process the selected subject
+    # Load content based on selection
     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, selected_subject + ".txt")
+        subject_file = os.path.join(essay_folder, f"{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 of notes and load vectorstore
+    # Display preview and create vectorstore
     if raw_text:
         st.subheader("Preview of Notes")
-        st.text_area("Preview Content:", value=raw_text[:1000], height=300, disabled=True)  # Display shorter preview
+        st.text_area("Preview Content:", value=raw_text[:1000], height=300, disabled=True)
 
-        # Preload vectorstore if not already cached
         if "vectorstore" not in st.session_state or st.session_state.vectorstore is None:
-            text_chunks = get_chunks(raw_text)
-            st.session_state.vectorstore = get_vectorstore(text_chunks)
+            chunks = get_chunks(raw_text)
+            st.session_state.vectorstore = get_vectorstore(chunks)
     else:
         st.warning("No content available for the selected subject.")
 
-    # Chat interface
+    # Question and response
     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.get("text", "No response found."))
+            st.write(response or "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()