Spaces:
Sleeping
Sleeping
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']) |