Spaces:
Build error
Build error
create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,108 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate
|
3 |
+
from llama_index.llms.huggingface import HuggingFaceInferenceAPI
|
4 |
+
from dotenv import load_dotenv
|
5 |
+
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
6 |
+
from llama_index.core import Settings
|
7 |
+
import os
|
8 |
+
import base64
|
9 |
+
|
10 |
+
# Load environment variables
|
11 |
+
load_dotenv()
|
12 |
+
|
13 |
+
# Configure the Llama index settings
|
14 |
+
Settings.llm = HuggingFaceInferenceAPI(
|
15 |
+
model_name="nltpt/Llama-3.2-3B-Instruct",
|
16 |
+
tokenizer_name="nltpt/Llama-3.2-3B-Instruct",
|
17 |
+
# model_name="google/gemma-1.1-7b-it",
|
18 |
+
# tokenizer_name="google/gemma-1.1-7b-it",
|
19 |
+
context_window=3000,
|
20 |
+
token=os.getenv("HF_TOKEN"),
|
21 |
+
max_new_tokens=512,
|
22 |
+
generate_kwargs={"temperature": 0.1},
|
23 |
+
)
|
24 |
+
Settings.embed_model = HuggingFaceEmbedding(
|
25 |
+
model_name="BAAI/bge-small-en-v1.5"
|
26 |
+
)
|
27 |
+
|
28 |
+
# Define the directory for persistent storage and data
|
29 |
+
PERSIST_DIR = "./db"
|
30 |
+
DATA_DIR = "data"
|
31 |
+
|
32 |
+
# Ensure data directory exists
|
33 |
+
os.makedirs(DATA_DIR, exist_ok=True)
|
34 |
+
os.makedirs(PERSIST_DIR, exist_ok=True)
|
35 |
+
|
36 |
+
def displayPDF(file):
|
37 |
+
with open(file, "rb") as f:
|
38 |
+
base64_pdf = base64.b64encode(f.read()).decode('utf-8')
|
39 |
+
pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600" type="application/pdf"></iframe>'
|
40 |
+
st.markdown(pdf_display, unsafe_allow_html=True)
|
41 |
+
|
42 |
+
def data_ingestion():
|
43 |
+
documents = SimpleDirectoryReader(DATA_DIR).load_data()
|
44 |
+
storage_context = StorageContext.from_defaults()
|
45 |
+
index = VectorStoreIndex.from_documents(documents)
|
46 |
+
index.storage_context.persist(persist_dir=PERSIST_DIR)
|
47 |
+
|
48 |
+
def handle_query(query):
|
49 |
+
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
|
50 |
+
index = load_index_from_storage(storage_context)
|
51 |
+
chat_text_qa_msgs = [
|
52 |
+
(
|
53 |
+
"user",
|
54 |
+
"""You are a Q&A assistant. Your main goal is to provide answers as accurately as possible, based on the context of the document provided. If the question does not match the context or is outside the scope of the document, advise the user to ask questions that are relevant to the document.
|
55 |
+
Context:
|
56 |
+
{context_str}
|
57 |
+
Question:
|
58 |
+
{query_str}
|
59 |
+
"""
|
60 |
+
)
|
61 |
+
]
|
62 |
+
|
63 |
+
text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
|
64 |
+
|
65 |
+
query_engine = index.as_query_engine(text_qa_template=text_qa_template)
|
66 |
+
answer = query_engine.query(query)
|
67 |
+
|
68 |
+
if hasattr(answer, 'response'):
|
69 |
+
return answer.response
|
70 |
+
elif isinstance(answer, dict) and 'response' in answer:
|
71 |
+
return answer['response']
|
72 |
+
else:
|
73 |
+
return "Sorry, I couldn't find an answer."
|
74 |
+
|
75 |
+
|
76 |
+
# Streamlit app initialization
|
77 |
+
st.title("(PDF) Information and Inference🗞️")
|
78 |
+
st.markdown("Retrieval-Augmented Generation")
|
79 |
+
st.markdown("Start chat ...🚀")
|
80 |
+
|
81 |
+
if 'messages' not in st.session_state:
|
82 |
+
st.session_state.messages = [{'role': 'assistant', "content": 'Hello! Upload PDF files and ask me anything about their content.'}]
|
83 |
+
|
84 |
+
with st.sidebar:
|
85 |
+
st.title("Menu:")
|
86 |
+
uploaded_files = st.file_uploader("Upload your PDF Files and Click on the Submit & Process Button", type="pdf", accept_multiple_files=True)
|
87 |
+
|
88 |
+
if st.button("Submit & Process"):
|
89 |
+
if uploaded_files:
|
90 |
+
with st.spinner("Processing..."):
|
91 |
+
for uploaded_file in uploaded_files:
|
92 |
+
filepath = os.path.join(DATA_DIR, uploaded_file.name)
|
93 |
+
with open(filepath, "wb") as f:
|
94 |
+
f.write(uploaded_file.getbuffer())
|
95 |
+
data_ingestion() # Process PDFs after they are uploaded
|
96 |
+
st.success("Done")
|
97 |
+
else:
|
98 |
+
st.warning("Please upload at least one PDF file.")
|
99 |
+
|
100 |
+
user_prompt = st.chat_input("Ask me anything about the content of the PDF(s):")
|
101 |
+
if user_prompt:
|
102 |
+
st.session_state.messages.append({'role': 'user', "content": user_prompt})
|
103 |
+
response = handle_query(user_prompt)
|
104 |
+
st.session_state.messages.append({'role': 'assistant', "content": response})
|
105 |
+
|
106 |
+
for message in st.session_state.messages:
|
107 |
+
with st.chat_message(message['role']):
|
108 |
+
st.write(message['content'])
|