import os import torch import streamlit as st from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_community.vectorstores import FAISS from langchain_core.prompts import PromptTemplate from langchain.chains import LLMChain from langchain_community.llms import HuggingFacePipeline from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM from dotenv import load_dotenv # Set Streamlit page configuration st.set_page_config(page_title="Chat with Notes and AI", page_icon=":books:", layout="wide") # Load environment variables load_dotenv() # Dolly-v2-3b model pipeline @st.cache_resource def load_pipeline(): 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.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 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 ) # Initialize Dolly pipeline generate_text = load_pipeline() # Create a 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}" ) # 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 def get_text_files_content(folder): text = "" for filename in os.listdir(folder): if filename.endswith('.txt'): with open(os.path.join(folder, filename), 'r', encoding='utf-8') as file: text += file.read() + "\n" return text # Converting text to 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 length_function=len ) chunks = text_splitter.split_text(raw_text) return chunks # Using Hugging Face embeddings model and FAISS to create vectorstore def get_vectorstore(chunks): embeddings = HuggingFaceEmbeddings( model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={'device': 'cpu'} # Ensure embeddings use CPU ) vectorstore = FAISS.from_texts(texts=chunks, embedding=embeddings) return vectorstore # Generating 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] if context: result_with_context = llm_context_chain.invoke({"instruction": question, "context": context}) return result_with_context # Fallback to instruction-only chain if no context is found return llm_chain.invoke({"instruction": question}) def main(): st.title("Chat with Notes :books:") # Initialize session state 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 # Content type selection content_type = st.sidebar.radio("Select Content Type:", ["Current Affairs", "Essays"]) # Handle Current Affairs (each subject has its own folder) 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) 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) # 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, 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 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 # 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.") # 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.get("text", "No response found.")) else: st.warning("Please load the content for the selected subject before asking a question.") if __name__ == '__main__': main()