Update app.py

app.py

@@ -15,56 +15,46 @@ st.set_page_config(page_title="Chat with Notes and AI", page_icon=":books:", lay
 # Load environment variables
 load_dotenv()
 
-# Dolly-v2-3b model pipeline
+# Optimized Dolly-v2 model pipeline
 @st.cache_resource
 def load_pipeline():
-    model_name = "databricks/dolly-v2-3b"
+    model_name = "databricks/dolly-v2-1b"  # Smaller model for CPU
 
-    # Load tokenizer
+    # Load tokenizer and model
     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.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)
+        torch_dtype=torch.float32,  # Use float32 for CPU
+        device_map="auto",
         trust_remote_code=True,
-        offload_folder="./offload_weights"  # Folder to store offloaded weights
+        offload_folder="./offload_weights"  # Folder to store weights if needed
     )
 
-    # Return text-generation pipeline
+    # Create text-generation pipeline
     return pipeline(
         task="text-generation",
         model=model,
         tokenizer=tokenizer,
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto",
-        return_full_text=True
+        max_new_tokens=50,  # Limit response length for speed
+        return_full_text=False,
+        device_map="auto"
     )
 
 # Initialize Dolly pipeline
 generate_text = load_pipeline()
 
-# Create a HuggingFace pipeline wrapper for LangChain
+# Create HuggingFace pipeline wrapper for LangChain
 hf_pipeline = HuggingFacePipeline(pipeline=generate_text)
 
-# 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}"
-)
+# Prompt templates
+prompt = PromptTemplate(input_variables=["instruction"], template="{instruction}")
+prompt_with_context = PromptTemplate(input_variables=["instruction", "context"], template="{instruction}\n\nInput:\n{context}")
 
 # Create LLM chains
 llm_chain = LLMChain(llm=hf_pipeline, prompt=prompt)
 llm_context_chain = LLMChain(llm=hf_pipeline, prompt=prompt_with_context)
 
-# Extracting text from .txt files
+# Extract text from .txt files
 def get_text_files_content(folder):
     text = ""
     for filename in os.listdir(folder):
@@ -73,36 +63,30 @@ def get_text_files_content(folder):
                 text += file.read() + "\n"
     return text
 
-# Converting text to chunks
+# Convert text into smaller chunks
 def get_chunks(raw_text):
     from langchain.text_splitter import CharacterTextSplitter
     text_splitter = CharacterTextSplitter(
         separator="\n",
-        chunk_size=1000,  # Reduced chunk size for faster processing
-        chunk_overlap=200,  # Smaller overlap for efficiency
+        chunk_size=512,  # Smaller chunks for faster processing
+        chunk_overlap=50,  # Minimal overlap
         length_function=len
     )
-    chunks = text_splitter.split_text(raw_text)
-    return chunks
+    return text_splitter.split_text(raw_text)
 
-# Using Hugging Face embeddings model and FAISS to create vectorstore
+# Create FAISS vectorstore
 def get_vectorstore(chunks):
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2",
-        model_kwargs={'device': 'cpu'}  # Ensure embeddings use CPU
+        model_kwargs={'device': 'cpu'}  # Force CPU usage for embeddings
     )
-    vectorstore = FAISS.from_texts(texts=chunks, embedding=embeddings)
-    return vectorstore
+    return FAISS.from_texts(texts=chunks, embedding=embeddings)
 
-# Generating response from user queries
+# Generate response from user queries
 def handle_question(question, vectorstore=None):
     if vectorstore:
-        # 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]
+        documents = vectorstore.similarity_search(question, k=1)  # Retrieve fewer chunks
+        context = "\n".join([doc.page_content for doc in documents])[:512]  # Shorter context
 
         if context:
             result_with_context = llm_context_chain.invoke({"instruction": question, "context": context})
@@ -122,26 +106,18 @@ def main():
     data_folder = "data"  # Current Affairs folders
     essay_folder = "essays"  # Essays folder
 
-    # Content type selection
+    # Sidebar for content selection
     content_type = st.sidebar.radio("Select Content Type:", ["Current Affairs", "Essays"])
 
-    # Handle Current Affairs (each subject has its own folder)
+    # Handle folder-based selection
     if content_type == "Current Affairs":
-        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)
+        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":
-        if os.path.exists(essay_folder):
-            subjects = [f.replace(".txt", "") for f in os.listdir(essay_folder) if f.endswith('.txt')]
-        else:
-            subjects = []
+        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 selected subject
+    # Process the selected subject
     raw_text = ""
     if content_type == "Current Affairs" and selected_subject:
         subject_folder = os.path.join(data_folder, selected_subject)
@@ -152,15 +128,15 @@ def main():
             with open(subject_file, "r", encoding="utf-8") as file:
                 raw_text = file.read()
 
-    # Display preview of notes
+    # Display preview of notes and load vectorstore
     if raw_text:
         st.subheader("Preview of Notes")
-        st.text_area("Preview Content:", value=raw_text[:2000], height=300, disabled=True)  # Show a snippet of the notes
+        st.text_area("Preview Content:", value=raw_text[:1000], height=300, disabled=True)  # Display shorter preview
 
-        # Create vectorstore for Current Affairs or Essays
-        text_chunks = get_chunks(raw_text)
-        vectorstore = get_vectorstore(text_chunks)
-        st.session_state.vectorstore = vectorstore
+        # 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)
     else:
         st.warning("No content available for the selected subject.")