Spaces:
Running
Running
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) |