Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
|
@@ -6,73 +6,180 @@ from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
|
|
| 6 |
import os
|
| 7 |
from dotenv import load_dotenv
|
| 8 |
import gradio as gr
|
|
|
|
|
|
|
| 9 |
|
| 10 |
# Load environment variables
|
| 11 |
load_dotenv()
|
| 12 |
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 19 |
parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
|
| 20 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 21 |
|
| 22 |
-
#
|
| 23 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 24 |
|
| 25 |
-
# Global variable to store documents loaded from user-uploaded files
|
| 26 |
-
vector_index = None
|
| 27 |
|
| 28 |
# File processing function
|
| 29 |
-
def load_files(file_path: str):
|
| 30 |
try:
|
| 31 |
global vector_index
|
| 32 |
document = SimpleDirectoryReader(input_files=[file_path], file_extractor=file_extractor).load_data()
|
|
|
|
| 33 |
vector_index = VectorStoreIndex.from_documents(document, embed_model=embed_model)
|
| 34 |
-
print(f"
|
| 35 |
filename = os.path.basename(file_path)
|
| 36 |
-
return f"Ready to give response on
|
| 37 |
except Exception as e:
|
| 38 |
-
return f"An error occurred {e}"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 39 |
|
|
|
|
|
|
|
| 40 |
def respond(message, history):
|
| 41 |
try:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 42 |
query_engine = vector_index.as_query_engine(llm=llm)
|
| 43 |
bot_message = query_engine.query(message)
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
# yield output
|
| 48 |
-
print(f"{datetime.now()}::message=>{str(bot_message)}")
|
| 49 |
-
return str(bot_message)
|
| 50 |
except Exception as e:
|
| 51 |
-
if e == "'NoneType' object has no attribute 'as_query_engine'":
|
| 52 |
-
return "upload file"
|
| 53 |
-
return f"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 54 |
|
| 55 |
# UI Setup
|
| 56 |
-
with gr.Blocks() as
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 77 |
if __name__ == "__main__":
|
| 78 |
-
|
|
|
|
| 6 |
import os
|
| 7 |
from dotenv import load_dotenv
|
| 8 |
import gradio as gr
|
| 9 |
+
import markdowm as md
|
| 10 |
+
import base64
|
| 11 |
|
| 12 |
# Load environment variables
|
| 13 |
load_dotenv()
|
| 14 |
|
| 15 |
+
llm_models = [
|
| 16 |
+
"mistralai/Mixtral-8x7B-Instruct-v0.1",
|
| 17 |
+
"meta-llama/Meta-Llama-3-8B-Instruct",
|
| 18 |
+
"mistralai/Mistral-7B-Instruct-v0.2",
|
| 19 |
+
"tiiuae/falcon-7b-instruct",
|
| 20 |
+
# "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
|
| 21 |
+
# "deepseek-ai/deepseek-vl2", ## 54GB > 10GB
|
| 22 |
+
# "deepseek-ai/deepseek-vl2-small", ## 32GB > 10GB
|
| 23 |
+
# "deepseek-ai/deepseek-vl2-tiny", ## high response time
|
| 24 |
+
# "deepseek-ai/deepseek-llm-7b-chat", ## 13GB > 10GB
|
| 25 |
+
# "deepseek-ai/deepseek-math-7b-instruct", ## 13GB > 10GB
|
| 26 |
+
# "deepseek-ai/deepseek-coder-33b-instruct", ## 66GB > 10GB
|
| 27 |
+
# "deepseek-ai/DeepSeek-R1-Zero", ## 688GB > 10GB
|
| 28 |
+
# "mistralai/Mixtral-8x22B-Instruct-v0.1", ## 281GB>10GB
|
| 29 |
+
# "NousResearch/Yarn-Mistral-7b-64k", ## 14GB>10GB
|
| 30 |
+
# "impira/layoutlm-document-qa", ## ERR
|
| 31 |
+
# "Qwen/Qwen1.5-7B", ## 15GB
|
| 32 |
+
# "Qwen/Qwen2.5-3B", ## high response time
|
| 33 |
+
# "google/gemma-2-2b-jpn-it", ## high response time
|
| 34 |
+
# "impira/layoutlm-invoices", ## bad req
|
| 35 |
+
# "google/pix2struct-docvqa-large", ## bad req
|
| 36 |
+
# "google/gemma-7b-it", ## 17GB > 10GB
|
| 37 |
+
# "google/gemma-2b-it", ## high response time
|
| 38 |
+
# "HuggingFaceH4/zephyr-7b-beta", ## high response time
|
| 39 |
+
# "HuggingFaceH4/zephyr-7b-gemma-v0.1", ## bad req
|
| 40 |
+
# "microsoft/phi-2", ## high response time
|
| 41 |
+
# "TinyLlama/TinyLlama-1.1B-Chat-v1.0", ## high response time
|
| 42 |
+
# "mosaicml/mpt-7b-instruct", ## 13GB>10GB
|
| 43 |
+
# "google/flan-t5-xxl" ## high respons time
|
| 44 |
+
# "NousResearch/Yarn-Mistral-7b-128k", ## 14GB>10GB
|
| 45 |
+
# "Qwen/Qwen2.5-7B-Instruct", ## 15GB>10GB
|
| 46 |
+
]
|
| 47 |
|
| 48 |
+
embed_models = [
|
| 49 |
+
"BAAI/bge-small-en-v1.5", # 33.4M
|
| 50 |
+
"NeuML/pubmedbert-base-embeddings",
|
| 51 |
+
"BAAI/llm-embedder", # 109M
|
| 52 |
+
"BAAI/bge-large-en" # 335M
|
| 53 |
+
]
|
| 54 |
+
|
| 55 |
+
# Global variable for selected model
|
| 56 |
+
selected_llm_model_name = llm_models[0] # Default to the first model in the list
|
| 57 |
+
selected_embed_model_name = embed_models[0] # Default to the first model in the list
|
| 58 |
+
vector_index = None
|
| 59 |
+
|
| 60 |
+
# Initialize the parser
|
| 61 |
parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
|
| 62 |
+
# Define file extractor with various common extensions
|
| 63 |
+
file_extractor = {
|
| 64 |
+
'.pdf': parser, # PDF documents
|
| 65 |
+
'.docx': parser, # Microsoft Word documents
|
| 66 |
+
'.doc': parser, # Older Microsoft Word documents
|
| 67 |
+
'.txt': parser, # Plain text files
|
| 68 |
+
'.csv': parser, # Comma-separated values files
|
| 69 |
+
'.xlsx': parser, # Microsoft Excel files (requires additional processing for tables)
|
| 70 |
+
'.pptx': parser, # Microsoft PowerPoint files (for slides)
|
| 71 |
+
'.html': parser, # HTML files (web pages)
|
| 72 |
+
# '.rtf': parser, # Rich Text Format files
|
| 73 |
+
# '.odt': parser, # OpenDocument Text files
|
| 74 |
+
# '.epub': parser, # ePub files (e-books)
|
| 75 |
|
| 76 |
+
# Image files for OCR processing
|
| 77 |
+
'.jpg': parser, # JPEG images
|
| 78 |
+
'.jpeg': parser, # JPEG images
|
| 79 |
+
'.png': parser, # PNG images
|
| 80 |
+
# '.bmp': parser, # Bitmap images
|
| 81 |
+
# '.tiff': parser, # TIFF images
|
| 82 |
+
# '.tif': parser, # TIFF images (alternative extension)
|
| 83 |
+
# '.gif': parser, # GIF images (can contain text)
|
| 84 |
+
|
| 85 |
+
# Scanned documents in image formats
|
| 86 |
+
'.webp': parser, # WebP images
|
| 87 |
+
'.svg': parser, # SVG files (vector format, may contain embedded text)
|
| 88 |
+
}
|
| 89 |
|
|
|
|
|
|
|
| 90 |
|
| 91 |
# File processing function
|
| 92 |
+
def load_files(file_path: str, embed_model_name: str):
|
| 93 |
try:
|
| 94 |
global vector_index
|
| 95 |
document = SimpleDirectoryReader(input_files=[file_path], file_extractor=file_extractor).load_data()
|
| 96 |
+
embed_model = HuggingFaceEmbedding(model_name=embed_model_name)
|
| 97 |
vector_index = VectorStoreIndex.from_documents(document, embed_model=embed_model)
|
| 98 |
+
print(f"Parsing done for {file_path}")
|
| 99 |
filename = os.path.basename(file_path)
|
| 100 |
+
return f"Ready to give response on {filename}"
|
| 101 |
except Exception as e:
|
| 102 |
+
return f"An error occurred: {e}"
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
# Function to handle the selected model from dropdown
|
| 106 |
+
def set_llm_model(selected_model):
|
| 107 |
+
global selected_llm_model_name
|
| 108 |
+
selected_llm_model_name = selected_model # Update the global variable
|
| 109 |
+
# print(f"Model selected: {selected_model_name}")
|
| 110 |
+
# return f"Model set to: {selected_model_name}"
|
| 111 |
|
| 112 |
+
|
| 113 |
+
# Respond function that uses the globally set selected model
|
| 114 |
def respond(message, history):
|
| 115 |
try:
|
| 116 |
+
# Initialize the LLM with the selected model
|
| 117 |
+
llm = HuggingFaceInferenceAPI(
|
| 118 |
+
model_name=selected_llm_model_name,
|
| 119 |
+
contextWindow=8192, # Context window size (typically max length of the model)
|
| 120 |
+
maxTokens=1024, # Tokens per response generation (512-1024 works well for detailed answers)
|
| 121 |
+
temperature=0.3, # Lower temperature for more focused answers (0.2-0.4 for factual info)
|
| 122 |
+
topP=0.9, # Top-p sampling to control diversity while retaining quality
|
| 123 |
+
frequencyPenalty=0.5, # Slight penalty to avoid repetition
|
| 124 |
+
presencePenalty=0.5, # Encourages exploration without digressing too much
|
| 125 |
+
token=os.getenv("TOKEN")
|
| 126 |
+
)
|
| 127 |
+
|
| 128 |
+
# Set up the query engine with the selected LLM
|
| 129 |
query_engine = vector_index.as_query_engine(llm=llm)
|
| 130 |
bot_message = query_engine.query(message)
|
| 131 |
+
|
| 132 |
+
print(f"\n{datetime.now()}:{selected_llm_model_name}:: {message} --> {str(bot_message)}\n")
|
| 133 |
+
return f"{selected_llm_model_name}:\n{str(bot_message)}"
|
|
|
|
|
|
|
|
|
|
| 134 |
except Exception as e:
|
| 135 |
+
if str(e) == "'NoneType' object has no attribute 'as_query_engine'":
|
| 136 |
+
return "Please upload a file."
|
| 137 |
+
return f"An error occurred: {e}"
|
| 138 |
+
|
| 139 |
+
def encode_image(image_path):
|
| 140 |
+
with open(image_path, "rb") as image_file:
|
| 141 |
+
return base64.b64encode(image_file.read()).decode('utf-8')
|
| 142 |
+
|
| 143 |
+
# Encode the images
|
| 144 |
+
github_logo_encoded = encode_image("Images/github-logo.png")
|
| 145 |
+
linkedin_logo_encoded = encode_image("Images/linkedin-logo.png")
|
| 146 |
+
website_logo_encoded = encode_image("Images/ai-logo.png")
|
| 147 |
|
| 148 |
# UI Setup
|
| 149 |
+
with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto Mono")]), css='footer {visibility: hidden}') as demo:
|
| 150 |
+
gr.Markdown("# DocBot📄🤖")
|
| 151 |
+
with gr.Tabs():
|
| 152 |
+
with gr.TabItem("Intro"):
|
| 153 |
+
gr.Markdown(md.description)
|
| 154 |
+
|
| 155 |
+
with gr.TabItem("DocBot"):
|
| 156 |
+
with gr.Accordion("=== IMPORTANT: READ ME FIRST ===", open=False):
|
| 157 |
+
guid = gr.Markdown(md.guide)
|
| 158 |
+
with gr.Row():
|
| 159 |
+
with gr.Column(scale=1):
|
| 160 |
+
file_input = gr.File(file_count="single", type='filepath', label="Step-1: Upload document")
|
| 161 |
+
# gr.Markdown("Dont know what to select check out in Intro tab")
|
| 162 |
+
embed_model_dropdown = gr.Dropdown(embed_models, label="Step-2: Select Embedding", interactive=True)
|
| 163 |
+
with gr.Row():
|
| 164 |
+
btn = gr.Button("Submit", variant='primary')
|
| 165 |
+
clear = gr.ClearButton()
|
| 166 |
+
output = gr.Text(label='Vector Index')
|
| 167 |
+
llm_model_dropdown = gr.Dropdown(llm_models, label="Step-3: Select LLM", interactive=True)
|
| 168 |
+
with gr.Column(scale=3):
|
| 169 |
+
gr.ChatInterface(
|
| 170 |
+
fn=respond,
|
| 171 |
+
chatbot=gr.Chatbot(height=500),
|
| 172 |
+
theme = "soft",
|
| 173 |
+
show_progress='full',
|
| 174 |
+
# cache_mode='lazy',
|
| 175 |
+
textbox=gr.Textbox(placeholder="Step-4: Ask me questions on the uploaded document!", container=False)
|
| 176 |
+
)
|
| 177 |
+
gr.HTML(md.footer.format(github_logo_encoded, linkedin_logo_encoded, website_logo_encoded))
|
| 178 |
+
# Set up Gradio interactions
|
| 179 |
+
llm_model_dropdown.change(fn=set_llm_model, inputs=llm_model_dropdown)
|
| 180 |
+
btn.click(fn=load_files, inputs=[file_input, embed_model_dropdown], outputs=output)
|
| 181 |
+
clear.click(lambda: [None] * 3, outputs=[file_input, embed_model_dropdown, output])
|
| 182 |
+
|
| 183 |
+
# Launch the demo with a public link option
|
| 184 |
if __name__ == "__main__":
|
| 185 |
+
demo.launch(share=True)
|