bug fix

app.py

@@ -61,13 +61,12 @@ def get_image_description(image):
     torch.cuda.empty_cache()
     gc.collect()
 
-    descriptions = []
     prompt = "[INST] <image>\nDescribe the image in a sentence [/INST]"
+    n = len(prompt)
 
     inputs = processor(prompt, image, return_tensors="pt").to("cuda:0")
     output = vision_model.generate(**inputs, max_new_tokens=100)
-    descriptions.append(processor.decode(output[0], skip_special_tokens=True))
-    return descriptions
+    return processor.decode(output[0][n:], skip_special_tokens=True)
 
 
 CSS = """
@@ -101,9 +100,14 @@ def get_vectordb(text, images, img_doc_files):
     descs = []
     for i in range(len(images)):
         try:
-            descs.append(img_doc_files[i]+"\n"+get_image_description(images[i])[0])
+            descs.append(img_doc_files[i] + "\n" + get_image_description(images[i]))
         except:
             descs.append("Could not generate image description due to some error")
+            print(img_doc_files)
+            print(images[i])
+            print()
+        print(descs[-1])
+        print()
 
     # image_descriptions = get_image_descriptions(images)
     image_dict = [{"image": image_to_bytes(img)} for img in images]
@@ -134,7 +138,7 @@ def get_vectordb(text, images, img_doc_files):
 def extract_only_text(reader):
     text = ""
     for _, page in enumerate(reader.pages):
-        text += page.extract_text()
+        text = page.extract_text()
     return text.strip()
 
 
@@ -152,7 +156,7 @@ def extract_data_from_pdfs(
     all_text = ""
 
     images = []
-    img_docs=[]
+    img_docs = []
     for doc in docs:
         if do_ocr == "Get Text With OCR":
             pdf_doc = DocumentFile.from_pdf(doc)
@@ -163,8 +167,9 @@ def extract_data_from_pdfs(
             all_text += extract_only_text(reader) + "\n\n"
 
         if include_images == "Include Images":
-            images.extend(extract_images([doc]))
-            img_docs.append(doc.split("/")[-1])
+            imgs = extract_images([doc])
+            images.extend(imgs)
+            img_docs.extend([doc.split("/")[-1] for _ in range(len(imgs))])
 
     progress(
         0.6, "Generating image descriptions and inserting everything into vectorDB"
@@ -266,13 +271,6 @@ def check_validity_and_llm(session_states):
     raise gr.Error("Please extract data first")
 
 
-def get_stats(vectordb):
-    eles = vectordb.get()
-    # words =
-    text_data = [f"Chunks: {len(eles)}", "HIII"]
-    return "\n".join(text_data), "", ""
-
-
 with gr.Blocks(css=CSS, theme=gr.themes.Soft(text_size=sizes.text_md)) as demo:
     vectordb = gr.State()
     doc_collection = gr.State(value=[])
@@ -462,10 +460,14 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft(text_size=sizes.text_md)) as demo:
         [vectordb, msg, num_context, img_context, chatbot, hf_token, model_path],
         [chatbot, references, ret_images],
     )
-    documents.change(lambda: "<h1 style='text-align: center'>Click the 'Extract' button to extract data from PDFs<h1>", None, prog)
+    documents.change(
+        lambda: "<h1 style='text-align: center'>Click the 'Extract' button to extract data from PDFs<h1>",
+        None,
+        prog,
+    )
 
     back_p1.click(lambda: gr.Tabs(selected=0), None, tabs)
 
     next_p1.click(check_validity_and_llm, session_states, tabs)
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()