Document-QA-bot / app.py
Sarath0x8f's picture
Update app.py
8ef6cb8 verified
raw
history blame
2.95 kB
from datetime import datetime
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_parse import LlamaParse
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
import os
from dotenv import load_dotenv
import gradio as gr
# Load environment variables
load_dotenv()
# Initialize the LLM and parser
llm = HuggingFaceInferenceAPI(
model_name="meta-llama/Meta-Llama-3-8B-Instruct",
token=os.getenv("TOKEN")
)
parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
file_extractor = {'.pdf': parser, '.docx': parser, '.doc': parser}
# Embedding model and index initialization (to be populated by uploaded files)
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
# Global variable to store documents loaded from user-uploaded files
vector_index = None
# File processing function
def load_files(file_path: str):
try:
global vector_index
document = SimpleDirectoryReader(input_files=[file_path], file_extractor=file_extractor).load_data()
vector_index = VectorStoreIndex.from_documents(document, embed_model=embed_model)
print(f"parsing done {file_path}")
filename = os.path.basename(file_path)
return f"Ready to give response on give {filename}"
except Exception as e:
return f"An error occurred {e}"
def respond(message, history):
try:
query_engine = vector_index.as_query_engine(llm=llm)
bot_message = query_engine.query(message)
# output = ""
# for chr in bot_message:
# output += chr
# yield output
print(f"{datetime.now()}::message=>{str(bot_message)}")
return str(bot_message)
except Exception as e:
if e == "'NoneType' object has no attribute 'as_query_engine'":
return "upload file"
return f"an error occurred {e}"
# UI Setup
with gr.Blocks() as demo:
with gr.Row():
with gr.Column(scale=1):
file_input = gr.File(file_count="single", type='filepath')
with gr.Column():
clear = gr.ClearButton()
btn = gr.Button("Submit", variant='primary')
output = gr.Text(label='Vector Index')
with gr.Column(scale=2):
gr.ChatInterface(fn=respond,
chatbot=gr.Chatbot(height=500),
textbox=gr.Textbox(placeholder="Ask me a yes or no question", container=False, scale=7),
examples=["summarize the document"]
)
# Action on button click to process file and load into index
btn.click(fn=load_files, inputs=file_input, outputs=output)
clear.click(lambda: [None]*2, outputs=[file_input, output])
# Launch the demo with public link option
if __name__ == "__main__":
demo.launch(share=True)