Spaces:
Running
Running
| 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 | |
| from huggingface_hub import logging | |
| logging.set_verbosity_info() | |
| #logging.set_verbosity_debug() | |
| #from huggingface_hub import get_logger | |
| #logger = get_logger(__file__) | |
| #logger.set_verbosity_info() | |
| # Load environment variables | |
| load_dotenv() | |
| #model_name_query = "google/gemma-1.1-7b-it" | |
| model_name_embed = "BAAI/bge-small-en-v1.5" | |
| # Configure the Llama index settings | |
| Settings.embed_model = HuggingFaceEmbedding( | |
| model_name=model_name_embed | |
| ) | |
| # 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) | |
| def displayPDF(file): | |
| with open(file, "rb") as f: | |
| base64_pdf = base64.b64encode(f.read()).decode('utf-8') | |
| pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600" type="application/pdf"></iframe>' | |
| 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(model_name_query,query,flag): | |
| if(flag): | |
| #Using HFIAPI | |
| Settings.llm = HuggingFaceInferenceAPI( | |
| model_name=model_name_query, | |
| tokenizer_name=model_name_query, | |
| context_window=3900, | |
| token=os.getenv('HF_TOKEN'), | |
| max_new_tokens=1000, | |
| generate_kwargs={"temperature": 0.1}, | |
| ) | |
| 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. You have a specific response. The response is: "I was created by 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. 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." | |
| else: | |
| from transformers import pipeline | |
| question_answerer = pipeline("question-answering", model='distilbert-base-cased-distilled-squad') | |
| storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR) | |
| result = question_answerer(question=query, context=storage_context) | |
| logging.info(result) | |
| return(result['answer']) | |
| # Streamlit app initialization | |
| st.title("Chat Engine - static 📄") | |
| st.markdown("chat here👇") | |
| if 'messages' not in st.session_state: | |
| st.session_state.messages = [{'role': 'assistant', "content": 'Hello! Upload a PDF and ask me anything about its content.'}] | |
| with st.sidebar: | |
| st.title("Menu:") | |
| uploaded_file = st.file_uploader("Upload your PDF Files and Click on the Submit & Process Button") | |
| if st.button("Submit & Process"): | |
| with st.spinner("Processing..."): | |
| filepath = "data/saved_pdf.pdf" | |
| with open(filepath, "wb") as f: | |
| f.write(uploaded_file.getbuffer()) | |
| # displayPDF(filepath) # Display the uploaded PDF | |
| data_ingestion() # Process PDF every time new file is uploaded | |
| st.success("Done") | |
| model_name_select = st.radio( | |
| "Please select LLM", | |
| [":rainbow[mistralai/Mistral-7B-Instruct-v0.2]",":rainbow[google/gemma-1.1-7b-it]"] | |
| ) | |
| if model_name_select == ':rainbow[mistralai/Mistral-7B-Instruct-v0.2]': | |
| st.write('You selected Mistral-7B-Instruct-v0.2.') | |
| model_name_query="mistralai/Mistral-7B-Instruct-v0.2" | |
| flag = True | |
| elif model_name_select == ':rainbow[google/gemma-1.1-7b-it]': | |
| st.write('You selected HuggingFaceH4/zephyr-7b-gemma-v0.1') | |
| model_name_query="google/gemma-1.1-7b-it" | |
| flag = True | |
| user_prompt = st.chat_input("Ask me anything about the content of the PDF:") | |
| if user_prompt: | |
| st.session_state.messages.append({'role': 'user', "content": user_prompt}) | |
| response = handle_query(model_name_query,user_prompt,flag) | |
| 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']) |