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) |