Basic retrieval and generation

app.py

@@ -1,7 +1,10 @@
 import io
+import os
 
 import streamlit as st
 import torch
+from byaldi import RAGMultiModalModel
+from pdf2image import convert_from_bytes
 from PIL import Image
 from transformers import (AutoModelForVision2Seq, AutoProcessor,
                           BitsAndBytesConfig)
@@ -9,9 +12,13 @@ from transformers.image_utils import load_image
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+@st.cache_resource  # Streamlit Caching decorator
+def load_model_embedding():
+    docs_retrieval_model = RAGMultiModalModel.from_pretrained("vidore/colsmolvlm-alpha")
+model_embedding = load_model_embedding()
 
 @st.cache_resource  # Streamlit Caching decorator
-def load_model():
+def load_model_vlm():
     checkpoint = "HuggingFaceTB/SmolVLM-Instruct"
     processor = AutoProcessor.from_pretrained(checkpoint)
     quantization_config = BitsAndBytesConfig(load_in_8bit=True)
@@ -20,51 +27,76 @@ def load_model():
         #torch_dtype=torch.bfloat16,
         quantization_config=quantization_config,
     )
-    return model
-model = load_model()
+    return model, processor
+model_vlm, processor_vlm = load_model_vlm()
 
 
-with st.sidebar:
-    "[Source Code](https://huggingface.co/spaces/deepakkarkala/multimodal-rag/tree/main)"
+
+def save_images_to_local(dataset, output_folder="data/"):
+    os.makedirs(output_folder, exist_ok=True)
+
+    for image_id, image in enumerate(dataset):
+        #if isinstance(image, str):
+        #    image = Image.open(image)
+
+        output_path = os.path.join(output_folder, f"image_{image_id}.png")
+        image.save(output_path, format="PNG")
+
 
 
 # Home page UI
+with st.sidebar:
+    "[Source Code](https://huggingface.co/spaces/deepakkarkala/multimodal-rag/tree/main)"
+
 st.title("📝 Image Q&A with VLM")
-uploaded_file = st.file_uploader("Upload an image", type=("png", "jpg"))
-question = st.text_input(
+uploaded_pdf = st.file_uploader("Upload PDF file", type=("pdf"))
+query = st.text_input(
     "Ask something about the image",
     placeholder="Can you describe me the image ?",
-    disabled=not uploaded_file,
+    disabled=not uploaded_pdf,
 )
 
+images = []
+if uploaded_pdf:
+    images = convert_from_bytes(uploaded_pdf.getvalue())
+    save_images_to_local(images)
+    # index documents using the document retrieval model
+    model_embedding.index(
+        input_path="data/", index_name="image_index", store_collection_with_index=False, overwrite=True
+    )
 
-if uploaded_file and question:
-    image_bytes = uploaded_file.read()
-    image = Image.open(io.BytesIO(image_bytes))
+
+
+if uploaded_pdf and query:
+    docs_retrieved = model_embedding.search(query, k=1)
+    image_similar_to_query = images[docs_retrieved[0]["doc_id"]]
 
     # Create input messages
     system_prompt = "You are an AI assistant. Your task is reply to user questions based on the provided image context."
-    messages = [
+    chat_template = [
         {"role": "system", "content": system_prompt},
         {
             "role": "user",
             "content": [
                 {"type": "image"},
-                {"type": "text", "text": question}
+                {"type": "text", "text": query}
             ]
         },
     ]
 
     # Prepare inputs
-    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-    inputs = processor(text=prompt, images=[image], return_tensors="pt")
+    prompt = processor_vlm.apply_chat_template(chat_template, add_generation_prompt=True)
+    inputs = processor_vlm(text=prompt, images=[image_similar_to_query], return_tensors="pt")
     inputs = inputs.to(DEVICE)
 
     # Generate outputs
-    generated_ids = model.generate(**inputs, max_new_tokens=500)
-    generated_texts = processor.batch_decode(
+    generated_ids = model_vlm.generate(**inputs, max_new_tokens=500)
+    #generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+
+    generated_texts = processor_vlm.batch_decode(
         generated_ids,
         skip_special_tokens=True,
+        clean_up_tokenization_spaces=False,
     )
     response = generated_texts[0]
     

requirements.txt

@@ -11,3 +11,6 @@ transformers
 accelerate>=0.26.0
 bitsandbytes
 pillow
+flash-attn
+byaldi
+pdf2image