import streamlit as st from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate from llama_index.llms.huggingface import HuggingFaceInferenceAPI from dotenv import load_dotenv from llama_index.embeddings.huggingface import HuggingFaceEmbedding from llama_index.core import Settings import os import base64 # Load environment variables load_dotenv() # Configure the Llama index settings Settings.llm = HuggingFaceInferenceAPI( model_name="google/gemma-1.1-7b-it", tokenizer_name="google/gemma-1.1-7b-it", context_window=3000, token=os.getenv("HF_TOKEN"), max_new_tokens=512, generate_kwargs={"temperature": 0.1}, ) Settings.embed_model = HuggingFaceEmbedding( model_name="BAAI/bge-small-en-v1.5" ) # Define the directory for persistent storage and data PERSIST_DIR = "./db" DATA_DIR = "data" # Ensure data directory exists os.makedirs(DATA_DIR, exist_ok=True) os.makedirs(PERSIST_DIR, exist_ok=True) # List preloaded PDFs PRELOADED_PDFS = [f for f in os.listdir(DATA_DIR) if f.endswith('.pdf')] def displayPDF(file): with open(file, "rb") as f: base64_pdf = base64.b64encode(f.read()).decode('utf-8') pdf_display = f'' st.markdown(pdf_display, unsafe_allow_html=True) def data_ingestion(): documents = SimpleDirectoryReader(DATA_DIR).load_data() storage_context = StorageContext.from_defaults() index = VectorStoreIndex.from_documents(documents) index.storage_context.persist(persist_dir=PERSIST_DIR) def handle_query(query): storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR) index = load_index_from_storage(storage_context) chat_text_qa_msgs = [ ( "user", """You are a Q&A assistant named CHATTO, created by Rishi. You have a specific response programmed for when users specifically ask about your creator, Rishi. The response is: "I was created by Rishi, an enthusiast in Artificial Intelligence. He is dedicated to solving complex problems and delivering innovative solutions. With a strong focus on machine learning, deep learning, Python, generative AI, NLP, and computer vision, Rishi is passionate about pushing the boundaries of AI to explore new possibilities." For all other inquiries, your main goal is to provide answers as accurately as possible, based on the instructions and context you have been given. If a question does not match the provided context or is outside the scope of the document, kindly advise the user to ask questions within the context of the document. Context: {context_str} Question: {query_str} """ ) ] text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs) query_engine = index.as_query_engine(text_qa_template=text_qa_template) answer = query_engine.query(query) if hasattr(answer, 'response'): return answer.response elif isinstance(answer, dict) and 'response' in answer: return answer['response'] else: return "Sorry, I couldn't find an answer." # Streamlit app initialization st.title("(PDF) Information and Inference") st.markdown("Retrieval-Augmented Generation") # Automatically load all PDFs from the data directory if 'pdfs_loaded' not in st.session_state: with st.spinner("Loading PDFs..."): data_ingestion() st.session_state.pdfs_loaded = True if 'messages' not in st.session_state: preloaded_pdfs_list = ", ".join(PRELOADED_PDFS) if PRELOADED_PDFS else "No PDFs available" st.session_state.messages = [{'role': 'assistant', "content": f'Hello! I have loaded the following PDFs: {preloaded_pdfs_list}. Ask me anything about their content.'}] user_prompt = st.chat_input("Ask me anything about the content of the PDFs:") if user_prompt: st.session_state.messages.append({'role': 'user', "content": user_prompt}) response = handle_query(user_prompt) st.session_state.messages.append({'role': 'assistant', "content": response}) for message in st.session_state.messages: with st.chat_message(message['role']): st.write(message['content'])