pratham0011
commited on
Commit
•
62f8d39
1
Parent(s):
a83a57d
Upload 5 files
Browse files- index.py +94 -0
- main.py +152 -0
- static/app.py +126 -0
- static/man-kddi.png +0 -0
- static/robot.png +0 -0
index.py
ADDED
@@ -0,0 +1,94 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI, UploadFile, File, HTTPException
|
2 |
+
from pydantic import BaseModel
|
3 |
+
from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate
|
4 |
+
from llama_index.llms.huggingface import HuggingFaceInferenceAPI
|
5 |
+
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
6 |
+
from llama_index.core import Settings
|
7 |
+
import os
|
8 |
+
from dotenv import load_dotenv
|
9 |
+
import shutil
|
10 |
+
|
11 |
+
# Load environment variables
|
12 |
+
load_dotenv()
|
13 |
+
|
14 |
+
app = FastAPI()
|
15 |
+
|
16 |
+
# Configure the Llama index settings
|
17 |
+
Settings.llm = HuggingFaceInferenceAPI(
|
18 |
+
model_name="meta-llama/Meta-Llama-3-8B-Instruct",
|
19 |
+
tokenizer_name="meta-llama/Meta-Llama-3-8B-Instruct",
|
20 |
+
context_window=3900,
|
21 |
+
token=os.getenv("HF_TOKEN"),
|
22 |
+
max_new_tokens=1000,
|
23 |
+
generate_kwargs={"temperature": 0.5},
|
24 |
+
)
|
25 |
+
Settings.embed_model = HuggingFaceEmbedding(
|
26 |
+
model_name="BAAI/bge-small-en-v1.5"
|
27 |
+
)
|
28 |
+
|
29 |
+
# Define the directory for persistent storage and data
|
30 |
+
PERSIST_DIR = "./db"
|
31 |
+
DATA_DIR = "data"
|
32 |
+
|
33 |
+
# Ensure data directory exists
|
34 |
+
os.makedirs(DATA_DIR, exist_ok=True)
|
35 |
+
os.makedirs(PERSIST_DIR, exist_ok=True)
|
36 |
+
|
37 |
+
class Query(BaseModel):
|
38 |
+
question: str
|
39 |
+
|
40 |
+
def data_ingestion():
|
41 |
+
documents = SimpleDirectoryReader(DATA_DIR).load_data()
|
42 |
+
storage_context = StorageContext.from_defaults()
|
43 |
+
index = VectorStoreIndex.from_documents(documents)
|
44 |
+
index.storage_context.persist(persist_dir=PERSIST_DIR)
|
45 |
+
|
46 |
+
def handle_query(query):
|
47 |
+
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
|
48 |
+
index = load_index_from_storage(storage_context)
|
49 |
+
chat_text_qa_msgs = [
|
50 |
+
(
|
51 |
+
"user",
|
52 |
+
"""You are Q&A assistant named CHAT-DOC. 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.
|
53 |
+
Context:
|
54 |
+
{context_str}
|
55 |
+
Question:
|
56 |
+
{query_str}
|
57 |
+
"""
|
58 |
+
)
|
59 |
+
]
|
60 |
+
text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
|
61 |
+
query_engine = index.as_query_engine(text_qa_template=text_qa_template)
|
62 |
+
answer = query_engine.query(query)
|
63 |
+
|
64 |
+
if hasattr(answer, 'response'):
|
65 |
+
return answer.response
|
66 |
+
elif isinstance(answer, dict) and 'response' in answer:
|
67 |
+
return answer['response']
|
68 |
+
else:
|
69 |
+
return "Sorry, I couldn't find an answer."
|
70 |
+
|
71 |
+
@app.post("/upload")
|
72 |
+
async def upload_file(file: UploadFile = File(...)):
|
73 |
+
file_extension = os.path.splitext(file.filename)[1].lower()
|
74 |
+
if file_extension not in [".pdf", ".docx", ".txt"]:
|
75 |
+
raise HTTPException(status_code=400, detail="Invalid file type. Only PDF, DOCX, and TXT are allowed.")
|
76 |
+
|
77 |
+
file_path = os.path.join(DATA_DIR, file.filename)
|
78 |
+
with open(file_path, "wb") as buffer:
|
79 |
+
shutil.copyfileobj(file.file, buffer)
|
80 |
+
|
81 |
+
data_ingestion()
|
82 |
+
return {"message": "File uploaded and processed successfully"}
|
83 |
+
|
84 |
+
@app.post("/query")
|
85 |
+
async def query_document(query: Query):
|
86 |
+
if not os.listdir(DATA_DIR):
|
87 |
+
raise HTTPException(status_code=400, detail="No document has been uploaded yet.")
|
88 |
+
|
89 |
+
response = handle_query(query.question)
|
90 |
+
return {"response": response}
|
91 |
+
|
92 |
+
if __name__ == "__main__":
|
93 |
+
import uvicorn
|
94 |
+
uvicorn.run(app, host="0.0.0.0", port=8000)
|
main.py
ADDED
@@ -0,0 +1,152 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# File: main.py
|
2 |
+
from fastapi import FastAPI, UploadFile, File, HTTPException
|
3 |
+
from pydantic import BaseModel
|
4 |
+
from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate
|
5 |
+
from llama_index.llms.huggingface import HuggingFaceInferenceAPI
|
6 |
+
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
|
7 |
+
from llama_index.core import Settings
|
8 |
+
import os
|
9 |
+
from dotenv import load_dotenv
|
10 |
+
import shutil
|
11 |
+
import uvicorn
|
12 |
+
import streamlit as st
|
13 |
+
import requests
|
14 |
+
import base64
|
15 |
+
import docx2txt
|
16 |
+
import threading
|
17 |
+
|
18 |
+
# Load environment variables
|
19 |
+
load_dotenv()
|
20 |
+
|
21 |
+
app = FastAPI()
|
22 |
+
|
23 |
+
# Configure the Llama index settings
|
24 |
+
Settings.llm = HuggingFaceInferenceAPI(
|
25 |
+
model_name="meta-llama/Meta-Llama-3-8B-Instruct",
|
26 |
+
tokenizer_name="meta-llama/Meta-Llama-3-8B-Instruct",
|
27 |
+
context_window=3900,
|
28 |
+
token=os.getenv("HF_TOKEN"),
|
29 |
+
max_new_tokens=1000,
|
30 |
+
generate_kwargs={"temperature": 0.5},
|
31 |
+
)
|
32 |
+
Settings.embed_model = HuggingFaceEmbedding(
|
33 |
+
model_name="BAAI/bge-small-en-v1.5"
|
34 |
+
)
|
35 |
+
|
36 |
+
# Define the directory for persistent storage and data
|
37 |
+
PERSIST_DIR = "./db"
|
38 |
+
DATA_DIR = "data"
|
39 |
+
|
40 |
+
# Ensure data directory exists
|
41 |
+
os.makedirs(DATA_DIR, exist_ok=True)
|
42 |
+
os.makedirs(PERSIST_DIR, exist_ok=True)
|
43 |
+
|
44 |
+
class Query(BaseModel):
|
45 |
+
question: str
|
46 |
+
|
47 |
+
def data_ingestion():
|
48 |
+
documents = SimpleDirectoryReader(DATA_DIR).load_data()
|
49 |
+
storage_context = StorageContext.from_defaults()
|
50 |
+
index = VectorStoreIndex.from_documents(documents)
|
51 |
+
index.storage_context.persist(persist_dir=PERSIST_DIR)
|
52 |
+
|
53 |
+
def handle_query(query):
|
54 |
+
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
|
55 |
+
index = load_index_from_storage(storage_context)
|
56 |
+
chat_text_qa_msgs = [
|
57 |
+
(
|
58 |
+
"user",
|
59 |
+
"""You are Q&A assistant named CHAT-DOC. 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.
|
60 |
+
Context:
|
61 |
+
{context_str}
|
62 |
+
Question:
|
63 |
+
{query_str}
|
64 |
+
"""
|
65 |
+
)
|
66 |
+
]
|
67 |
+
text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
|
68 |
+
query_engine = index.as_query_engine(text_qa_template=text_qa_template)
|
69 |
+
answer = query_engine.query(query)
|
70 |
+
|
71 |
+
if hasattr(answer, 'response'):
|
72 |
+
return answer.response
|
73 |
+
elif isinstance(answer, dict) and 'response' in answer:
|
74 |
+
return answer['response']
|
75 |
+
else:
|
76 |
+
return "Sorry, I couldn't find an answer."
|
77 |
+
|
78 |
+
@app.post("/upload")
|
79 |
+
async def upload_file(file: UploadFile = File(...)):
|
80 |
+
file_extension = os.path.splitext(file.filename)[1].lower()
|
81 |
+
if file_extension not in [".pdf", ".docx", ".txt"]:
|
82 |
+
raise HTTPException(status_code=400, detail="Invalid file type. Only PDF, DOCX, and TXT are allowed.")
|
83 |
+
|
84 |
+
file_path = os.path.join(DATA_DIR, file.filename)
|
85 |
+
with open(file_path, "wb") as buffer:
|
86 |
+
shutil.copyfileobj(file.file, buffer)
|
87 |
+
|
88 |
+
data_ingestion()
|
89 |
+
return {"message": "File uploaded and processed successfully"}
|
90 |
+
|
91 |
+
@app.post("/query")
|
92 |
+
async def query_document(query: Query):
|
93 |
+
if not os.listdir(DATA_DIR):
|
94 |
+
raise HTTPException(status_code=400, detail="No document has been uploaded yet.")
|
95 |
+
|
96 |
+
response = handle_query(query.question)
|
97 |
+
return {"response": response}
|
98 |
+
|
99 |
+
# Streamlit UI
|
100 |
+
def streamlit_ui():
|
101 |
+
st.title("Chat with your Document 📄")
|
102 |
+
st.markdown("Chat here👇")
|
103 |
+
|
104 |
+
icons = {"assistant": "🤖", "user": "👤"}
|
105 |
+
|
106 |
+
if 'messages' not in st.session_state:
|
107 |
+
st.session_state.messages = [{'role': 'assistant', "content": 'Hello! Upload a PDF, DOCX, or TXT file and ask me anything about its content.'}]
|
108 |
+
|
109 |
+
for message in st.session_state.messages:
|
110 |
+
with st.chat_message(message['role'], avatar=icons[message['role']]):
|
111 |
+
st.write(message['content'])
|
112 |
+
|
113 |
+
with st.sidebar:
|
114 |
+
st.title("Menu:")
|
115 |
+
uploaded_file = st.file_uploader("Upload your document (PDF, DOCX, TXT)", type=["pdf", "docx", "txt"])
|
116 |
+
if st.button("Submit & Process") and uploaded_file:
|
117 |
+
with st.spinner("Processing..."):
|
118 |
+
files = {"file": (uploaded_file.name, uploaded_file.getvalue(), uploaded_file.type)}
|
119 |
+
response = requests.post("http://localhost:8000/upload", files=files)
|
120 |
+
if response.status_code == 200:
|
121 |
+
st.success("File uploaded and processed successfully")
|
122 |
+
else:
|
123 |
+
st.error("Error uploading file")
|
124 |
+
|
125 |
+
user_prompt = st.chat_input("Ask me anything about the content of the document:")
|
126 |
+
|
127 |
+
if user_prompt:
|
128 |
+
st.session_state.messages.append({'role': 'user', "content": user_prompt})
|
129 |
+
with st.chat_message("user", avatar=icons["user"]):
|
130 |
+
st.write(user_prompt)
|
131 |
+
|
132 |
+
# Trigger assistant's response retrieval and update UI
|
133 |
+
with st.spinner("Thinking..."):
|
134 |
+
response = requests.post("http://localhost:8000/query", json={"question": user_prompt})
|
135 |
+
if response.status_code == 200:
|
136 |
+
assistant_response = response.json()["response"]
|
137 |
+
with st.chat_message("assistant", avatar=icons["assistant"]):
|
138 |
+
st.write(assistant_response)
|
139 |
+
st.session_state.messages.append({'role': 'assistant', "content": assistant_response})
|
140 |
+
else:
|
141 |
+
st.error("Error querying document")
|
142 |
+
|
143 |
+
def run_fastapi():
|
144 |
+
uvicorn.run(app, host="0.0.0.0", port=8000)
|
145 |
+
|
146 |
+
if __name__ == "__main__":
|
147 |
+
# Start FastAPI in a separate thread
|
148 |
+
fastapi_thread = threading.Thread(target=run_fastapi)
|
149 |
+
fastapi_thread.start()
|
150 |
+
|
151 |
+
# Run Streamlit (this will run in the main thread)
|
152 |
+
streamlit_ui()
|
static/app.py
ADDED
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
import docx2txt
|
10 |
+
|
11 |
+
# Load environment variables
|
12 |
+
load_dotenv()
|
13 |
+
|
14 |
+
icons = {"assistant": "robot.png", "user": "man-kddi.png"}
|
15 |
+
|
16 |
+
# Configure the Llama index settings
|
17 |
+
Settings.llm = HuggingFaceInferenceAPI(
|
18 |
+
model_name="meta-llama/Meta-Llama-3-8B-Instruct",
|
19 |
+
tokenizer_name="meta-llama/Meta-Llama-3-8B-Instruct",
|
20 |
+
context_window=3900,
|
21 |
+
token=os.getenv("HF_TOKEN"),
|
22 |
+
max_new_tokens=1000,
|
23 |
+
generate_kwargs={"temperature": 0.5},
|
24 |
+
)
|
25 |
+
Settings.embed_model = HuggingFaceEmbedding(
|
26 |
+
model_name="BAAI/bge-small-en-v1.5"
|
27 |
+
)
|
28 |
+
|
29 |
+
# Define the directory for persistent storage and data
|
30 |
+
PERSIST_DIR = "./db"
|
31 |
+
DATA_DIR = "data"
|
32 |
+
|
33 |
+
# Ensure data directory exists
|
34 |
+
os.makedirs(DATA_DIR, exist_ok=True)
|
35 |
+
os.makedirs(PERSIST_DIR, exist_ok=True)
|
36 |
+
|
37 |
+
def displayPDF(file):
|
38 |
+
with open(file, "rb") as f:
|
39 |
+
base64_pdf = base64.b64encode(f.read()).decode('utf-8')
|
40 |
+
pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600" type="application/pdf"></iframe>'
|
41 |
+
st.markdown(pdf_display, unsafe_allow_html=True)
|
42 |
+
|
43 |
+
def displayDOCX(file):
|
44 |
+
text = docx2txt.process(file)
|
45 |
+
st.text_area("Document Content", text, height=400)
|
46 |
+
|
47 |
+
def displayTXT(file):
|
48 |
+
with open(file, "r") as f:
|
49 |
+
text = f.read()
|
50 |
+
st.text_area("Document Content", text, height=400)
|
51 |
+
|
52 |
+
def data_ingestion():
|
53 |
+
documents = SimpleDirectoryReader(DATA_DIR).load_data()
|
54 |
+
storage_context = StorageContext.from_defaults()
|
55 |
+
index = VectorStoreIndex.from_documents(documents)
|
56 |
+
index.storage_context.persist(persist_dir=PERSIST_DIR)
|
57 |
+
|
58 |
+
def handle_query(query):
|
59 |
+
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
|
60 |
+
index = load_index_from_storage(storage_context)
|
61 |
+
chat_text_qa_msgs = [
|
62 |
+
(
|
63 |
+
"user",
|
64 |
+
"""You are Q&A assistant named CHAT-DOC. 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.
|
65 |
+
Context:
|
66 |
+
{context_str}
|
67 |
+
Question:
|
68 |
+
{query_str}
|
69 |
+
"""
|
70 |
+
)
|
71 |
+
]
|
72 |
+
text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
|
73 |
+
query_engine = index.as_query_engine(text_qa_template=text_qa_template)
|
74 |
+
answer = query_engine.query(query)
|
75 |
+
|
76 |
+
if hasattr(answer, 'response'):
|
77 |
+
return answer.response
|
78 |
+
elif isinstance(answer, dict) and 'response' in answer:
|
79 |
+
return answer['response']
|
80 |
+
else:
|
81 |
+
return "Sorry, I couldn't find an answer."
|
82 |
+
|
83 |
+
# Streamlit app initialization
|
84 |
+
st.title("Chat with your Document 📄")
|
85 |
+
st.markdown("Chat here👇")
|
86 |
+
|
87 |
+
if 'messages' not in st.session_state:
|
88 |
+
st.session_state.messages = [{'role': 'assistant', "content": 'Hello! Upload a PDF, DOCX, or TXT file and ask me anything about its content.'}]
|
89 |
+
|
90 |
+
for message in st.session_state.messages:
|
91 |
+
with st.chat_message(message['role'], avatar=icons[message['role']]):
|
92 |
+
st.write(message['content'])
|
93 |
+
|
94 |
+
with st.sidebar:
|
95 |
+
st.title("Menu:")
|
96 |
+
uploaded_file = st.file_uploader("Upload your document (PDF, DOCX, TXT)", type=["pdf", "docx", "txt"])
|
97 |
+
if st.button("Submit & Process") and uploaded_file:
|
98 |
+
with st.spinner("Processing..."):
|
99 |
+
file_extension = os.path.splitext(uploaded_file.name)[1].lower()
|
100 |
+
filepath = os.path.join(DATA_DIR, "uploaded_file" + file_extension)
|
101 |
+
with open(filepath, "wb") as f:
|
102 |
+
f.write(uploaded_file.getbuffer())
|
103 |
+
|
104 |
+
if file_extension == ".pdf":
|
105 |
+
displayPDF(filepath)
|
106 |
+
elif file_extension == ".docx":
|
107 |
+
displayDOCX(filepath)
|
108 |
+
elif file_extension == ".txt":
|
109 |
+
displayTXT(filepath)
|
110 |
+
|
111 |
+
data_ingestion() # Process file every time a new file is uploaded
|
112 |
+
st.success("Done")
|
113 |
+
|
114 |
+
user_prompt = st.chat_input("Ask me anything about the content of the document:")
|
115 |
+
|
116 |
+
if user_prompt and uploaded_file:
|
117 |
+
st.session_state.messages.append({'role': 'user', "content": user_prompt})
|
118 |
+
with st.chat_message("user", avatar=icons["user"]):
|
119 |
+
st.write(user_prompt)
|
120 |
+
|
121 |
+
# Trigger assistant's response retrieval and update UI
|
122 |
+
with st.spinner("Thinking..."):
|
123 |
+
response = handle_query(user_prompt)
|
124 |
+
with st.chat_message("assistant", avatar=icons["assistant"]):
|
125 |
+
st.write(response)
|
126 |
+
st.session_state.messages.append({'role': 'assistant', "content": response})
|
static/man-kddi.png
ADDED
static/robot.png
ADDED